Diff Coverage

            for model_param, main_param in zip(model_group, main_group):
                model_gradient = model_param.main_grad
                if model_gradient is None:
                    model_gradient = model_param.grad
                main_param.grad = model_gradient
                model_param.grad = None

        for fp32_group in self.fp32_from_fp32_groups:
            for model_param in fp32_group:

        for fp32_group in self.fp32_from_fp32_groups:
            for model_param in fp32_group:
                model_gradient = model_param.main_grad
                if model_gradient is not None:
                    model_param.grad = model_gradient

    @torch.no_grad()
    def _copy_main_params_to_model_params(self) -> None:
        """Copy updated fp32 main params to fp16/bf16 model params."""

            # The final dtype is applied inside the atomic build (05 stage-5
            # item 5); the Trainer no longer patches it afterwards.
            model_init_dtype=self.config.model_init_dtype,
        )
        _validate_optimize_dtype(
            self.model,
            fp32_main_params=self.config.optimizer.fp32_main_params,
        )
        self.model_config = self.model.config

        """Step optimizers and schedulers after gradient accumulation."""
        optimizers = self.optimizer if isinstance(self.optimizer, list) else [self.optimizer]
        for optimizer in optimizers:
            # Keep zeros_like dispatched so lazily created optimizer states retain DTensor metadata.
            with SkipDTensorDispatch(no_skip={torch.zeros_like}):
                optimizer.step()
            optimizer.zero_grad()

        schedulers = (