Fix: some validation

Author: S1ro1

examples/fsdp2/nd_parallel.py

@@ -77,17 +77,17 @@ def forward(model, batch, optimizer, accelerator: Accelerator):
         buffers=[input_ids, labels], buffer_seq_dims=[1, 1], no_restore_buffers={input_ids, labels}
     ):
         # To get the proper loss value, we need to average across devices that are participating in data parallel/context parallel training
-        # loss_reduce_grp = (
-        #     accelerator.torch_device_mesh["dp_cp"].get_group()
-        #     if accelerator.parallelism_config.dp_cp_dim_names
-        #     else None
-        # )
+        loss_reduce_grp = (
+            accelerator.torch_device_mesh["dp_cp"].get_group()
+            if accelerator.parallelism_config.dp_cp_dim_names
+            else None
+        )
         outputs = model(**batch)
         loss = outputs.loss
         accelerator.backward(loss)
         optimizer.step()
         optimizer.zero_grad()
-        # dist.all_reduce(loss, op=dist.ReduceOp.AVG, group=loss_reduce_grp)
+        dist.all_reduce(loss, op=dist.ReduceOp.AVG, group=loss_reduce_grp)
 
     return loss
 
@@ -134,7 +134,8 @@ def train(args):
     dataloader = DataLoader(dataset, batch_size=1, collate_fn=create_collate_fn())
 
     model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)
-    model = fix_model(model)
+    if parallelism_config.cp_enabled:
+        model = fix_model(model)
 
     total_num_steps = min(args.num_steps, len(dataloader))
     performance_tracker = PerformanceTracker(warmup_steps=5)

src/accelerate/accelerator.py

@@ -448,6 +448,12 @@ def __init__(
 
         parallelism_config = self._setup_parallelism_config(parallelism_config, torch_tp_plugin)
 
+        # TODO: Siro - figure out a better place where this can go (needs to be above AcceleratorState init)
+        if parallelism_config and parallelism_config.cp_enabled and fsdp_plugin is None:
+            raise ValueError(
+                "`cp_enabled` is set to `True` in the `parallelism_config`, but no `fsdp_plugin` was provided. We need a `fsdp_plugin` to use `cp_enabled=True`, as we also shard the model across the device mesh to save more memory"
+            )
+
         kwargs = self.init_handler.to_kwargs() if self.init_handler is not None else {}
         kwargs["parallelism_config"] = parallelism_config
         self.state = AcceleratorState(