FSDP2 Not Compatible w/ HSDP (DeviceMesh Ignored By State)

[schopra8]

** Environment **

---------------------------------                                                                                                                                                                                                                                            
System Environment Report                                                                                                                                                                                                                                                    
Created: 2025-07-28 20:34:47 UTC                                                                                                                                                                                                                                             
---------------------------------                                                                                                                                                                                                                                            
                                                                                                                                                                                                                                                                             
PyTorch information                                                                                                                                                                                                                                                          
-------------------                                                                                                                                                                                                                                                          
PyTorch version: 2.7.1+cu128                                                                                                                                                                                                                                                 
Is debug build: False                                                                                                                                                                                                                                                        
CUDA used to build PyTorch: 12.8                                                                                                                                                                                                                                             
ROCM used to build PyTorch: N/A                                                                                                                                                                                                                                              
                                                                                                                                                                                                                                                                             
OS: Ubuntu 22.04.4 LTS (x86_64)                                                                                                                                                                                                                                              
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0                                                                                                                                                                                                                           
Clang version: Could not collect                                                                                                                                                                                                                                             
CMake version: Could not collect                                                                                                                                                                                                                                             
Libc version: glibc-2.35                                                                                                                                                                                                                                                     
                                                                                                                                                                                                                                                                             
Python version: 3.10.12 (main, Mar 22 2024, 16:50:05) [GCC 11.4.0] (64-bit runtime)                                                                                                                                                                                          
Python platform: Linux-6.5.0-1023-aws-x86_64-with-glibc2.35                                                                                                                                                                                                                  
Is CUDA available: True                                                                                                                                                                                                                                                      
CUDA runtime version: Could not collect                                                                                                                                                                                                                                      
CUDA_MODULE_LOADING set to: LAZY                                                                                                                                                                                                                                             
GPU models and configuration:                                                                                                                                                                                                                                                
GPU 0: NVIDIA H100 80GB HBM3                                                                                                                                                                                                                                                 
GPU 1: NVIDIA H100 80GB HBM3                                                                                                                                                                                                                                                 
GPU 2: NVIDIA H100 80GB HBM3                                                                                                                                                                                                                                                 
GPU 3: NVIDIA H100 80GB HBM3                                                                                                                                                                                                                                                 
GPU 4: NVIDIA H100 80GB HBM3                                                                                                                                                                                                                                                 
GPU 5: NVIDIA H100 80GB HBM3                                                                                                                                                                                                                                                 
GPU 6: NVIDIA H100 80GB HBM3                                                                                                                                                                                                                                                 
GPU 7: NVIDIA H100 80GB HBM3 

Nvidia driver version: 535.183.01                                                                                                                                                                                                                                            
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                       x86_64
CPU op-mode(s):                     32-bit, 64-bit
Address sizes:                      48 bits physical, 57 bits virtual
Byte Order:                         Little Endian
CPU(s):                             104
On-line CPU(s) list:                0-103
Vendor ID:                          GenuineIntel
Model name:                         Intel Xeon Processor (SapphireRapids)
CPU family:                         6
Model:                              143
Thread(s) per core:                 1
Core(s) per socket:                 104
Socket(s):                          1
Stepping:                           4
BogoMIPS:                           5600.00
Flags:                              fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology cpuid tsc_known_freq pni pclmulqdq vmx ssse3 fma cx16 pcid sse4
_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch cpuid_fault invpcid_single ssbd ibrs ibpb stibp ibrs_enhanced tpr_shadow flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid av
x512f avx512dq rdseed adx smap avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves avx_vnni avx512_bf16 wbnoinvd arat vnmi avx512vbmi umip pku ospke waitpkg avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopc
ntdq la57 rdpid bus_lock_detect cldemote movdiri movdir64b fsrm md_clear serialize tsxldtrk avx512_fp16 arch_capabilities
Virtualization:                     VT-x
Hypervisor vendor:                  KVM
Virtualization type:                full
L1d cache:                          3.3 MiB (104 instances)
L1i cache:                          3.3 MiB (104 instances)
L2 cache:                           416 MiB (104 instances)
L3 cache:                           16 MiB (1 instance)
NUMA node(s):                       1
NUMA node0 CPU(s):                  0-103
Vulnerability Gather data sampling: Not affected
Vulnerability Itlb multihit:        Not affected
Vulnerability L1tf:                 Not affected
Vulnerability Mds:                  Not affected
Vulnerability Meltdown:             Not affected
Vulnerability Mmio stale data:      Unknown: No mitigations
Vulnerability Retbleed:             Not affected
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass:    Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:           Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:           Mitigation; Enhanced / Automatic IBRS; IBPB conditional; RSB filling; PBRSB-eIBRS SW sequence; BHI Syscall hardening, KVM SW loop
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Mitigation; TSX disabled

Versions of relevant libraries:
[pip3] flake8==7.3.0
[pip3] mypy==1.16.1
[pip3] mypy_extensions==1.1.0
[pip3] numpy==2.1.1
[pip3] nvidia-cublas-cu12==12.8.3.14
[pip3] nvidia-cuda-cupti-cu12==12.8.57
[pip3] nvidia-cuda-nvrtc-cu12==12.8.61
[pip3] nvidia-cuda-runtime-cu12==12.8.57
[pip3] nvidia-cudnn-cu12==9.7.1.26
[pip3] nvidia-cufft-cu12==11.3.3.41
[pip3] nvidia-curand-cu12==10.3.9.55
[pip3] nvidia-cusolver-cu12==11.7.2.55
[pip3] nvidia-cusparse-cu12==12.5.7.53
[pip3] nvidia-cusparselt-cu12==0.6.3
[pip3] nvidia-nccl-cu12==2.26.2
[pip3] nvidia-nvjitlink-cu12==12.8.61

pip3] nvidia-cuda-runtime-cu12==12.8.57
[pip3] nvidia-cudnn-cu12==9.7.1.26
[pip3] nvidia-cufft-cu12==11.3.3.41
[pip3] nvidia-curand-cu12==10.3.9.55
[pip3] nvidia-cusolver-cu12==11.7.2.55
[pip3] nvidia-cusparse-cu12==12.5.7.53
[pip3] nvidia-cusparselt-cu12==0.6.3
[pip3] nvidia-nccl-cu12==2.26.2
[pip3] nvidia-nvjitlink-cu12==12.8.61
[pip3] nvidia-nvtx-cu12==12.8.55
[pip3] pytorch-lightning==2.5.2
[pip3] pytorch-ranger==0.1.1
[pip3] torch==2.7.1+cu128
[pip3] torch-optimizer==0.3.0
[pip3] torchaudio==2.7.1+cu128
[pip3] torchmetrics==1.7.1
[pip3] torchvision==0.22.1+cu128
[pip3] triton==3.3.1
[conda] Could not collect


Composer information
--------------------
Composer Version: 0.32.0
Composer Commit Hash: None
CPU Model: Intel Xeon Processor (SapphireRapids)
CPU Count: 104
Number of Nodes: 1
GPU Model: NVIDIA H100 80GB HBM3
GPUs per Node: 1
GPU Count: 1
CUDA Device Count: 8

** To reproduce

Steps to reproduce the behavior:

1. Train with this parallelism_config

  parallelism_config:
    fsdp: # FSPD2
      device_mesh:
        - [0, 1, 2, 3]
        - [4, 5, 6, 7]
      reshard_after_forward: False # Zero-2 (Gradients + Optimizer Sharding)
      activation_checkpointing: True # Activation Checkpointing
      load_monolith_rank0_only: True  # only rank0 touches the file
      state_dict_type: 'full'         # we are loading/saving full ckpts

2. Put print statements in

   composer/composer/core/state.py

   Line 524 in 4ae29b1

   self.device_mesh: Optional[DeviceMesh] = _create_device_mesh(self.device, self.fsdp_config, self.tp_config)

   to log the create device mesh. The device mesh will be flat (1-D) not (2-D).

Expected behavior

Per the documentation for FSDP2 in Composer, we should be able to provide a DeviceMesh in the FSDP2Config. When I provide a 2D configuration (as seen above), this is ignored.

If you read the code in state.py, the device_mesh property of FSDP2Config is never used. Instead, the HSDP (2D) config is built if fsdp_config.data_parallel_replicate_degree is set. In FSDP this should work. But, in FSDP2 these properties are hardcoded and are not settable.

As a result, it is not possible to use FSDP2 with HSDP (2D) configs.

Additional context

[schopra8]

Two hacks to get this working for now.

Monkey Patch FSDP2Config

                # Override the FSDP2Config properties so that we can get this to work ...
                if parallelism_config['fsdp'].get('data_parallel_replicate_degree',
                                                  None) is not None:
                    FSDP2Config.data_parallel_replicate_degree = property(
                        lambda self: parallelism_config['fsdp'].data_parallel_replicate_degree)
                if parallelism_config['fsdp'].get('data_parallel_shard_degree', None) is not None:
                    FSDP2Config.data_parallel_shard_degree = property(
                        lambda self: parallelism_config['fsdp'].data_parallel_shard_degree)
                parallelism_config_args['fsdp2'] = FSDP2Config.from_compatible_attrs(
                    parallelism_config['fsdp'])

You cannot set data_parallel_shard_degree OR data_parallel_replicate_degree because they are properties of the dataclass. A real fix will require changing the DataClass to make these configurable attributes (not read-only properties).

Force a Synchronization in state.py `_create_device_mesh
If you don't do this, you'll often hit a race condition and training WILL NOT start in HSDP settings with FSDP2.

def _create_device_mesh(
    device: Device,
    fsdp_config: Optional[FSDPConfig | FSDP2Config],
    tp_config: Optional[TPConfig],
) -> Optional[DeviceMesh]:
    if fsdp_config is None:
        return None

    # Gather dimensions and names for the device mesh
    dims: list[int] = []
    names: list[str] = []
    if fsdp_config.data_parallel_replicate_degree is not None:
        dims.append(fsdp_config.data_parallel_replicate_degree)
        names.append(ParallelismType.DATA_PARALLEL_REPLICATE.value)
    dims.append(fsdp_config.data_parallel_shard_degree)
    names.append(ParallelismType.DATA_PARALLEL_SHARD.value)
    if tp_config is not None:
        dims.append(tp_config.tensor_parallel_degree)
        names.append(ParallelismType.TENSOR_PARALLEL.value)

    # Fill in the unspecified dimensions
    product_of_dims = 1
    unspecified_dim_names = []
    for dim, name in zip(dims, names):
        if dim != -1:
            product_of_dims *= dim
        else:
            unspecified_dim_names.append(name)
    if len(unspecified_dim_names) > 1:
        raise ValueError(
            f'Found multiple parallelism dimensions with -1: {unspecified_dim_names}. '
            'Only one is allowed, which is set to fill the remaining dimensions.',
        )
    elif len(unspecified_dim_names) == 1:
        if product_of_dims > dist.get_world_size():
            raise ValueError(
                f'World size {dist.get_world_size()} is greater than the product of the specified parallelism degrees '
                f'{product_of_dims}. Please ensure the product of the specified parallelism degrees matches the world ',
                f'size. Currently specified degrees are {names=}, {dims=}. One dimension can also be left as -1, which '
                'will automatically be specified to ensure the product matches the world size.',
            )
        remaining_dimension = dist.get_world_size() // product_of_dims
        if remaining_dimension * product_of_dims != dist.get_world_size():
            raise ValueError(
                f'World size {dist.get_world_size()} is not divisible by the product of the specified '
                'parallelism degrees. Please ensure the product of the specified parallelism degrees '
                'matches the world size.',
            )
        for i, dim in enumerate(dims):
            if dim == -1:
                dims[i] = remaining_dimension
                log.info(f'Automatically setting {names[i]} to have parallelization degree {remaining_dimension}.')
                break
    else:
        if product_of_dims != dist.get_world_size():
            raise ValueError(
                f'World size {dist.get_world_size()} does not equal the product of the specified parallelism degrees '
                f'{product_of_dims}. Please ensure the product of the specified parallelism degrees matches the world ',
                f'size. Currently specified degrees are {names=}, {dims=}. One dimension can also be left as -1, which '
                'will automatically be specified to ensure the product matches the world size.',
            )

    device_type = device.name
    if device_type == 'gpu':
        device_type = 'cuda'

    # --------------------------------------------------------------------------------------------
    # New Code
    # --------------------------------------------------------------------------------------------
    device_mesh =  init_device_mesh(device_type=device_type, mesh_shape=tuple(dims), mesh_dim_names=tuple(names))

    # Force initialization of the FSDP2 process groups
    if device_mesh is not None and isinstance(fsdp_config, FSDP2Config):
        # Extract the FSDP portion of the mesh
        if device_mesh.mesh_dim_names is not None and ParallelismType.DATA_PARALLEL_REPLICATE.value in device_mesh.mesh_dim_names:
            # This is HSDP - initialize process groups
            dummy = device.tensor_to_device(torch.zeros(1))
            
            # Get the FSDP sub-mesh
            fsdp_mesh_dim_names = []
            if ParallelismType.DATA_PARALLEL_REPLICATE.value in device_mesh.mesh_dim_names:
                fsdp_mesh_dim_names.append(ParallelismType.DATA_PARALLEL_REPLICATE.value)
            fsdp_mesh_dim_names.append(ParallelismType.DATA_PARALLEL_SHARD.value)
            fsdp_mesh = device_mesh[tuple(fsdp_mesh_dim_names)]
            
            # Initialize both dimensions
            print("About to do all reduce ....")
            for dim in range(fsdp_mesh.ndim):
                pg = fsdp_mesh.get_group(dim)
                dist.all_reduce(dummy, group=pg)
            
            print("About to synchronize ....")
            if torch.cuda.is_available():
                torch.cuda.synchronize()
            dist.barrier()

    print("Sync complete ....")
    return device_mesh