diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py index 031b56a5c26..6c907bc25ec 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py @@ -2499,7 +2499,14 @@ def _init_each_parameter_group_buffers(self, meta_device_init_fp8_params): gc.collect() torch.cuda.empty_cache() - m.to_empty(device=self.device, recurse=False) + # Materialize only meta tensors in a module, preserving + # non-meta tensors that are already initialized on device. + m._apply( + lambda t: ( + torch.empty_like(t, device=self.device) if t.is_meta else t + ), + recurse=False, + ) if ( HAVE_TE and is_te_min_version("0.9.0")