def get_fsdp_submesh_2d(device_mesh: DeviceMesh):
    """Get the submesh for FSDP.

    Args:
        device_mesh (DeviceMesh): The full device mesh.

    Returns:
        DeviceMesh: The submesh for FSDP.
    """
    if device_mesh.mesh.ndim == 2:
        submesh = device_mesh['weight_parallel']
    elif device_mesh.mesh.ndim == 3:
        raise RuntimeError(f'HSDP + MoE is not supported.')
    else:
        raise ValueError(f'{device_mesh.mesh.ndim=} not supported for MoE.')

    return submesh
Get the submesh for FSDP.

    Args:
        device_mesh (DeviceMesh): The full device mesh.

    Returns:
        DeviceMesh: The submesh for FSDP.
    
llmfoundry/models/layers/ffn.py