[compile] Fix graphbreaks in moe split; scale_grad

ghstack-source-id: a4efe51
Pull Request resolved: #2771

Author: IvanKobzarev

recipes/full_finetune_distributed.py

@@ -342,6 +342,9 @@ def setup(self, cfg: DictConfig) -> None:
             self._compile_loss = compile.get("loss", True)
             self._compile_optimizer_step = compile.get("optimizer_step", False)
             self._compile_scale_grads = compile.get("scale_grads", True)
+        if self._compile_model:
+            # Enable capture_scalar_outputs to compile non-grouped-mm path of moe, that uses split
+            torch._dynamo.config.capture_scalar_outputs = True
 
         # This indirection is needed to apply torch.compile to scale_grads step.
         self._grad_scaler = training.scale_grads_
@@ -941,7 +944,7 @@ def train(self) -> None:
 
                         # Manually scale the gradients from unnormalized loss by total # of tokens
                         self._grad_scaler(
-                            self._model.parameters(),
+                            list(self._model.parameters()),
                             self.world_size / num_tokens,
                             False if self.parallel_dims.tp_enabled else None,
                         )