DCP Checkpoint Load Fails for _extra_state when training in FSDP2

[ajWithNucleus]

Describe the bug

DCP Checkpoint Load Fails for _extra_state when training in FSDP2

Steps/Code to reproduce bug

Save a DCP checkpoint and try loading it back when using FSDP

[rank4]: ValueError: Size mismatch between saved torch.Size([2322]) and current: torch.Size([4]) for model.diffusion_trans
former.layers.0.feed_forward.ffn._extra_state

Expected behavior

We should be able to load back the model properly to resume training.

Environment overview (please complete the following information)

• Environment location: Baremetal
• Method of Transformer Engine install: pip install, v2.3
• If method of install is [Docker], provide docker pull & docker run commands used

Environment details

If NVIDIA docker image is used you don't need to specify these.
Otherwise, please provide:

• OS version: Ubuntu 22.04
• PyTorch version: 2.7
• Python version: 3.12
• Transformer Engine version: v2.3
• CUDA version: 12.8
• CUDNN version: ~9

Device details

• GPU model: H100s

Additional context

Add any other context about the problem here.

[ajWithNucleus]

cc: @ksivaman

[ajWithNucleus]

Diving deeper, I found multiple issue that can be resolved by the following,

When loading the, directly initializing model and trying to load the checkpoint results in _extra_state being set None which when serialized results in 4 bytes. In general extra state gets initialized in prepare_forward method, so getting extra_state before forward results in None

ValueError: Size mismatch between saved torch.Size([2322]) and current: torch.Size([4]) for model.diffusion_trans
former.layers.0.feed_forward.ffn._extra_state

This can be mitigated by forcing the model to set extra state properly after initialization

        with te.fp8_autocast(enabled=args.model.diffusion_model.use_fp8_ffn, fp8_recipe=fp8_recipe):
            for module in model.modules():
                if isinstance(module, te.LayerNormMLP):
                    module.init_fp8_metadata(2) #LayerNormMLP requires num_gemms to be set to 2

After adding this, we are still get the below error

 ValueError: Size mismatch between saved torch.Size([2322]) and current: torch.Size([2318]) for model.diffusion_trans
former.layers.0.feed_forward.ffn._extra_state

We are still short by 4 bytes, when I dig deeper, the issue turns out to be an additional serialization in get_extra_state method in base.py, ref:

PyTorch can handle dict returns in get_state fns, when we manually serialize and save to disk there are additional 4bytes which don't match when we directly compare the return of get_state. We can remove this additional manual serialization as pyTorch natively handles this

        torch.cuda.synchronize()
        # state_serialized = bytearray(pickle.dumps(state))
        # state_serialized = torch.frombuffer(state_serialized, dtype=torch.uint8)
        return state

Now since we are returning dict, we can facilitate set_extra_state to simply return the dict it gets as is:

   def set_extra_state(self, state: torch.Tensor) -> None:
        """Load previous state."""
        # Load state
        if isinstance(state, torch.Tensor):
            # No FP8 is indicated by an empty tensor we don't need to unpickle.
            if state.numel() == 0:
                return
            # Default format: byte tensor with pickled data
            state = pickle.loads(state.detach().cpu().numpy().tobytes())
        elif isinstance(state, io.BytesIO):
            # Deprecated format with io.BytesIO
            state.seek(0)
            state = torch.load(state, map_location="cuda")
        elif isinstance(state, dict):
            return state
        else:
            raise RuntimeError("Unsupported checkpoint format.")

        if state is None:
            return

This should allow us to load /resume saved checkpoint.

@ksivaman is there a reason why we manually serialize extra_state instead of relying on PyTorch's machinery ?

If not we can, I can raise a PR with fix and update the documentation on how to load/resume checkpoint?