diff --git a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py index a852ef226b4..2b02b1a9f57 100644 --- a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py +++ b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py @@ -68,6 +68,7 @@ class FullyShardedDataParallel(_BaseDataParallel): _MODULE_TYPE_REGISTRY: Dict[str, set] = { "column": { "ColumnParallelLinear", + "LinearCrossEntropyModule", "TEColumnParallelLinear", "TELayerNormColumnParallelLinear", "TEColumnParallelGroupedLinear",