From 79fc9564ea5991b25e9c657dcab0d82babd66945 Mon Sep 17 00:00:00 2001 From: YannAhlgrim Date: Sat, 16 May 2026 11:52:10 +0200 Subject: [PATCH] DDP guard if single GPUgio --- src/train.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/src/train.py b/src/train.py index 19dde24..29f0722 100644 --- a/src/train.py +++ b/src/train.py @@ -214,11 +214,16 @@ def main(args, resume_preempt=False): ipe_scale=ipe_scale, use_bfloat16=use_bfloat16, ) - encoder = DistributedDataParallel(encoder, static_graph=True) - predictor = DistributedDataParallel(predictor, static_graph=True) - target_encoder = DistributedDataParallel(target_encoder) - for p in target_encoder.parameters(): - p.requires_grad = False + if dist.is_available() and dist.is_initialized() and world_size > 1: + encoder = DistributedDataParallel(encoder, static_graph=True) + predictor = DistributedDataParallel(predictor, static_graph=True) + target_encoder = DistributedDataParallel(target_encoder) + 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_scheduler = (