remove nccl stuff

PrimeIntellect-ai · Oct 7, 2024 · ce84987 · ce84987
1 parent 38b5aea
commit ce84987
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/zeroband/train.py b/src/zeroband/train.py
@@ -152,7 +152,7 @@ def train(config: Config):
         num = 1 if isinstance(config.train.ac_ckpt, bool) else config.train.ac_ckpt
         apply_ac_ckpt(model, num)
 
-    elastic_device_mesh = ElasticDeviceMesh("nccl")
+    elastic_device_mesh = ElasticDeviceMesh()
 
     mp_policy = MixedPrecisionPolicy(
         param_dtype=torch.bfloat16, reduce_dtype=torch.float32 if config.train.reduce_fp32 else None