def dtensorify_param(
    param: nn.Parameter,
    mesh: DeviceMesh,
    placements: list[Placement],
):
    """Construct a DTensor from an already sharded local parameter."""
    param_dtensor = DTensor.from_local(
        param.data,
        device_mesh=mesh,
        placements=placements,
        run_check=False,
    )
    return nn.Parameter(param_dtensor)
Construct a DTensor from an already sharded local parameter.
llmfoundry/models/layers/ffn.py