DDP guard if single GPUgio

This commit is contained in:
YannAhlgrim
2026-05-16 11:52:10 +02:00
parent bad00f306e
commit 79fc9564ea
+10 -5
View File
@@ -214,11 +214,16 @@ def main(args, resume_preempt=False):
ipe_scale=ipe_scale, ipe_scale=ipe_scale,
use_bfloat16=use_bfloat16, use_bfloat16=use_bfloat16,
) )
encoder = DistributedDataParallel(encoder, static_graph=True) if dist.is_available() and dist.is_initialized() and world_size > 1:
predictor = DistributedDataParallel(predictor, static_graph=True) encoder = DistributedDataParallel(encoder, static_graph=True)
target_encoder = DistributedDataParallel(target_encoder) predictor = DistributedDataParallel(predictor, static_graph=True)
for p in target_encoder.parameters(): target_encoder = DistributedDataParallel(target_encoder)
p.requires_grad = False for p in target_encoder.parameters():
p.requires_grad = False
else:
logger.info("Distributed process group not initialized; running without DDP")
for p in target_encoder.parameters():
p.requires_grad = False
# -- momentum schedule # -- momentum schedule
momentum_scheduler = ( momentum_scheduler = (