Initialize torch.distributed.
(args)
| 210 | |
| 211 | |
| 212 | def initialize_distributed(args): |
| 213 | """Initialize torch.distributed.""" |
| 214 | if torch.distributed.is_initialized(): |
| 215 | if mpu.model_parallel_is_initialized(): |
| 216 | if args.model_parallel_size != mpu.get_model_parallel_world_size(): |
| 217 | raise ValueError( |
| 218 | "model_parallel_size is inconsistent with prior configuration." |
| 219 | "We currently do not support changing model_parallel_size." |
| 220 | ) |
| 221 | return False |
| 222 | else: |
| 223 | if args.model_parallel_size > 1: |
| 224 | warnings.warn( |
| 225 | "model_parallel_size > 1 but torch.distributed is not initialized via SAT." |
| 226 | "Please carefully make sure the correctness on your own." |
| 227 | ) |
| 228 | mpu.initialize_model_parallel(args.model_parallel_size) |
| 229 | return True |
| 230 | # the automatic assignment of devices has been moved to arguments.py |
| 231 | if args.device == "cpu": |
| 232 | pass |
| 233 | else: |
| 234 | torch.cuda.set_device(args.device) |
| 235 | # Call the init process |
| 236 | init_method = "tcp://" |
| 237 | args.master_ip = os.getenv("MASTER_ADDR", "localhost") |
| 238 | |
| 239 | if args.world_size == 1: |
| 240 | from sat.helpers import get_free_port |
| 241 | |
| 242 | default_master_port = str(get_free_port()) |
| 243 | else: |
| 244 | default_master_port = "6000" |
| 245 | args.master_port = os.getenv("MASTER_PORT", default_master_port) |
| 246 | init_method += args.master_ip + ":" + args.master_port |
| 247 | torch.distributed.init_process_group( |
| 248 | backend=args.distributed_backend, world_size=args.world_size, rank=args.rank, init_method=init_method |
| 249 | ) |
| 250 | |
| 251 | # Set the model-parallel / data-parallel communicators. |
| 252 | mpu.initialize_model_parallel(args.model_parallel_size) |
| 253 | |
| 254 | # Set vae context parallel group equal to model parallel group |
| 255 | from sgm.util import set_context_parallel_group, initialize_context_parallel |
| 256 | |
| 257 | if args.model_parallel_size <= 2: |
| 258 | set_context_parallel_group(args.model_parallel_size, mpu.get_model_parallel_group()) |
| 259 | else: |
| 260 | initialize_context_parallel(2) |
| 261 | # mpu.initialize_model_parallel(1) |
| 262 | # Optional DeepSpeed Activation Checkpointing Features |
| 263 | if args.deepspeed: |
| 264 | import deepspeed |
| 265 | |
| 266 | deepspeed.init_distributed( |
| 267 | dist_backend=args.distributed_backend, world_size=args.world_size, rank=args.rank, init_method=init_method |
| 268 | ) |
| 269 | # # It seems that it has no negative influence to configure it even without using checkpointing. |
no test coverage detected