def set_ffn_device_mesh(
    ffn: nn.Module,
    moe_world_size: int,
    device_mesh: DeviceMesh,
    get_fsdp_submesh: Callable[[DeviceMesh], DeviceMesh],
):
    """Sets the device mesh in FSDP kwargs.

    Args:
        ffn (nn.Module): The FFN module.
        moe_world_size (int): The MoE world size.
        device_mesh (DeviceMesh): The full device mesh.
        get_fsdp_submesh (Callable[[DeviceMesh], DeviceMesh]): A function to get the fsdp submesh.

    Raises:
        RuntimeError: If the device mesh is 3D.
        ValueError: If the device mesh is not 2D or 3D.
    """
    if moe_world_size > 1:
        expert_mesh = device_mesh['expert_parallel']
        expert_placements: list[Placement] = [Shard(0)]
        # Register in two loops as you cannot overwrite parameters while iterating over named_parameters()
        dtensorified_params = [
            (
                name,
                dtensorify_param(
                    param=parameter,
                    mesh=expert_mesh,
                    placements=expert_placements,
                ),
            ) for name, parameter in \
                ffn.experts.mlp.named_parameters()  # type: ignore
        ]
        for name, dtensorified_param in dtensorified_params:
            ffn.experts.mlp.register_parameter(  # type: ignore
                name,
                dtensorified_param,
            )

        submesh = get_fsdp_submesh(device_mesh)

        ffn.experts._fsdp_kwargs_dict = {  # type: ignore
            'device_mesh': submesh,
        }
Sets the device mesh in FSDP kwargs.

    Args:
        ffn (nn.Module): The FFN module.
        moe_world_size (int): The MoE world size.
        device_mesh (DeviceMesh): The full device mesh.
        get_fsdp_submesh (Callable[[DeviceMesh], DeviceMesh]): A function to get the fsdp submesh.

    Raises:
        RuntimeError: If the device mesh is 3D.
        ValueError: If the device mesh is not 2D or 3D.
    
llmfoundry/models/layers/ffn.py