Add fused_adam, quantized_model_init, and fsdp2 example

[pstjohn]

Summary

• Fix FusedAdam to work with PyTorch-native FSDP2 (fully_shard) when parameters are DTensor-wrapped Float8Tensor/QuantizedTensor
• Fix fuse_wgrad_accumulation guard to avoid crashing with vanilla FSDP2 (previously assumed Megatron-Core FSDP exclusively)
• Add examples for quantized_model_init on single-GPU (main.py) and multi-GPU FSDP2 (fully_shard.py)

Note: fuse_wgrad_accumulation remains incompatible with vanilla FSDP2

fuse_wgrad_accumulation still cannot be used with vanilla FSDP2. The feature writes weight gradients directly into main_grad and returns None to autograd, bypassing FSDP2's reduce-scatter. Each rank ends up with an unreduced gradient. Megatron-Core FSDP solves this by wiring get_main_grad() into its own reduce-scatter infrastructure. Vanilla FSDP2 does not yet expose an equivalent hook.

Fixes #2682

[greptile-apps]

Greptile Summary

This PR enables FusedAdam optimizer to work correctly with PyTorch-native FSDP2 (fully_shard) when parameters are DTensor-wrapped Float8Tensor/QuantizedTensor objects. The core fix extracts local tensors from DTensor wrappers before passing them to multi-tensor kernels and state initialization routines.

Key changes:

• Modified FusedAdam to extract _local_tensor from DTensor params/gradients before operations
• Added comprehensive test suite for FSDP2 + FusedAdam integration across multiple FP8 recipes
• Includes tests for checkpoint save/load, safetensors export, and store_param_remainders mode
• Added well-documented examples demonstrating single-GPU and multi-GPU FSDP2 usage with quantized_model_init
• Correctly marks fuse_wgrad_accumulation + FSDP2 test as xfail (incompatible without Megatron-Core hooks)

Confidence Score: 5/5

• This PR is safe to merge with minimal risk
• The changes are well-isolated to DTensor handling in FusedAdam, thoroughly tested with 8 new test cases covering edge cases, and include comprehensive documentation. The xfail test properly documents the known limitation with fuse_wgrad_accumulation.
• No files require special attention

Important Files Changed

Filename	Overview
transformer_engine/pytorch/optimizers/fused_adam.py	Adds DTensor support for FSDP2 by extracting local tensors before operations, ensuring optimizer states are plain CUDA tensors
tests/pytorch/distributed/run_fsdp2_fused_adam.py	Comprehensive test suite covering FusedAdam with FSDP2, FP8 params, checkpoint save/load, and expected failure for fuse_wgrad_accumulation
examples/pytorch/quantized_model_init/fully_shard.py	Multi-GPU FSDP2 example showing meta-device init, sharding, FP8 quantization, and checkpoint/safetensors export

Sequence Diagram

sequenceDiagram
    participant User
    participant Model as TransformerLayer<br/>(FP8 params via quantized_model_init)
    participant FSDP2 as fully_shard<br/>(DTensor wrapping)
    participant Optimizer as FusedAdam<br/>(FP32 master weights)
    participant Kernel as Multi-tensor kernels

    User->>Model: Build model in quantized_model_init context
    Model-->>User: Float8Tensor/QuantizedTensor params
    User->>FSDP2: Apply fully_shard(model)
    FSDP2-->>User: DTensor-wrapped params
    Note over FSDP2: DTensor._local_tensor = Float8Tensor
    User->>Optimizer: FusedAdam(params, master_weights=True)
    Optimizer->>Optimizer: Extract local_tensor from DTensor
    Optimizer->>Optimizer: Dequantize QuantizedTensor
    Optimizer-->>User: FP32 master_param states initialized
    User->>Model: Forward + backward pass
    Model-->>Optimizer: Gradients (DTensor-wrapped)
    Optimizer->>Optimizer: Extract p_grad._local_tensor
    Optimizer->>Optimizer: Extract p._local_tensor for FP8 params
    Optimizer->>Kernel: Call multi_tensor_adam with plain CUDA tensors
    Kernel-->>Optimizer: Updated FP32 master weights
    Optimizer->>Optimizer: Quantize and update FP8 params

Loading

_{Last reviewed commit: 96c123e}

[greptile-apps]

_{6 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

_{6 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[cspades]

LGTM, clean edits.

[cspades]

Should we use dequantize(dtype=torch.float32), to fuse the cast into the de-quantization's output buffer? (Likely not a big deal since I don't think this will change anything numerically, and you only call this function during init and whenever you save and load DCP checkpoints.)

[greptile-apps]

_{6 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

_{6 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[XueSongTap]

@pstjohn Hi, thanks for the great work! Does this PR plan to also handle the BF16 path? I noticed the BF16 branch still operates on the original p/p_grad without unwrapping when they're DTensors. In my experiments with FSDP2 + BF16, I'm seeing non-trivial overhead during the optimizer step from repeated DTensor dispatch. Curious if that's intentional or a planned follow-up.