HY / trainer /models /loader /fsdp_load.py

Upload batch 65: 500 files (0.01 GiB)

74da989 verified about 2 months ago

16.3 kB

# SPDX-License-Identifier: Apache-2.0

# Adapted from torchtune
# Copyright 2024 The TorchTune Authors.
# Copyright 2025 The Trainer Authors.

import contextlib
from collections.abc import Callable, Generator
from itertools import chain
from typing import Any
import os

import torch
import torch.distributed as dist
from torch import nn
from torch.distributed import DeviceMesh, init_device_mesh
from torch.distributed._tensor import distribute_tensor
from torch.distributed.fsdp import (CPUOffloadPolicy, FSDPModule,
                                    MixedPrecisionPolicy, fully_shard)
from torch.nn.modules.module import _IncompatibleKeys
from torch.nn.parallel import DistributedDataParallel as DDP

from functools import partial
from torch.distributed.fsdp import FullStateDictConfig, FullyShardedDataParallel as FSDP, MixedPrecision, ShardingStrategy, StateDictType
from torch.distributed.fsdp.api import CPUOffload
from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy, transformer_auto_wrap_policy

from trainer.distributed import get_world_group
from trainer.logger import init_logger
from trainer.models.loader.utils import (get_param_names_mapping,
                                           hf_to_custom_state_dict)
from trainer.models.loader.weight_utils import safetensors_weights_iterator
from trainer.utils import set_mixed_precision_policy

logger = init_logger(__name__)


# TODO(PY): move this to utils elsewhere
@contextlib.contextmanager
def set_default_dtype(dtype: torch.dtype) -> Generator[None, None, None]:
    """
    Context manager to set torch's default dtype.
    Args:
        dtype (torch.dtype): The desired default dtype inside the context manager.
    Returns:
        ContextManager: context manager for setting default dtype.
    Example:
        >>> with set_default_dtype(torch.bfloat16):
        >>>     x = torch.tensor([1, 2, 3])
        >>>     x.dtype
        torch.bfloat16
    """
    old_dtype = torch.get_default_dtype()
    torch.set_default_dtype(dtype)
    try:
        yield
    finally:
        torch.set_default_dtype(old_dtype)


# TODO(PY): add compile option
def maybe_load_fsdp_model(
    load_from_dir: str,
    ar_action_load_from_dir: str,
    cls_name: str,
    model_cls: type[nn.Module],
    init_params: dict[str, Any],
    weight_dir_list: list[str],
    device: torch.device,
    hsdp_replicate_dim: int,
    hsdp_shard_dim: int,
    param_dtype: torch.dtype,
    reduce_dtype: torch.dtype,
    cpu_offload: bool = False,
    fsdp_inference: bool = False,
    output_dtype: torch.dtype | None = None,
    training_mode: bool = True,
    pin_cpu_memory: bool = True,
) -> torch.nn.Module:
    """
    Load the model with FSDP if is training, else load the model without FSDP.
    """
    # NOTE(will): cast_forward_inputs=True shouldn't be needed as we are
    # manually casting the inputs to the model
    mp_policy = MixedPrecisionPolicy(param_dtype,
                                     reduce_dtype,
                                     output_dtype,
                                     cast_forward_inputs=False)

    set_mixed_precision_policy(
        param_dtype=param_dtype,
        reduce_dtype=reduce_dtype,
        output_dtype=output_dtype,
        mp_policy=mp_policy,
    )

    # with set_default_dtype(param_dtype), torch.device("meta"):
    with set_default_dtype(param_dtype):
        logger.info(f"loading from: {load_from_dir}")
        logger.info(f"model:", model_cls)
        # model = model_cls(**init_params)
        model = model_cls.from_pretrained(
            load_from_dir, local_attn_size=-1, sink_size=0)

        
        if "Action" in cls_name:
            model.add_discrete_action_parameters()
            if ar_action_load_from_dir is not None:
                from safetensors.torch import load_file
                state_dict = load_file(ar_action_load_from_dir)
                model.load_state_dict(state_dict, strict=True)
            logger.info(f"loading from: {ar_action_load_from_dir}")
            # model.add_channel_concat_parameters()

        model.to(torch.float32)

    # Check if we should use FSDP
    use_fsdp = training_mode or fsdp_inference

    # Disable FSDP for MPS as it's not compatible
    from trainer.platforms import current_platform
    if current_platform.is_mps():
        use_fsdp = False
        logger.info("Disabling FSDP for MPS platform as it's not compatible")

    if use_fsdp:
        world_size = hsdp_replicate_dim * hsdp_shard_dim
        if not training_mode and not fsdp_inference:
            hsdp_replicate_dim = world_size
            hsdp_shard_dim = 1

        device_mesh = init_device_mesh(
            "cuda",
            # (Replicate(), Shard(dim=0))
            mesh_shape=(hsdp_replicate_dim, hsdp_shard_dim),
            mesh_dim_names=("replicate", "shard"),
        )
        shard_model(model,
                    cpu_offload=cpu_offload,
                    reshard_after_forward=True,
                    mp_policy=mp_policy,
                    mesh=device_mesh,
                    fsdp_shard_conditions=model._fsdp_shard_conditions,
                    pin_cpu_memory=pin_cpu_memory)

    for n, p in chain(model.named_parameters(), model.named_buffers()):
        if p.is_meta:
            raise RuntimeError(
                f"Unexpected param or buffer {n} on meta device.")
        # Avoid unintended computation graph accumulation during inference
        if isinstance(p, torch.nn.Parameter):
            p.requires_grad = False
    return model

def build_2d_groups(rank: int, world_size: int, fsdp_group_size: int):
    assert world_size % fsdp_group_size == 0, "world_size 必须能整除 FSDP_GROUP_SIZE"
    num_replicas = world_size // fsdp_group_size

    shard_group_id = rank // fsdp_group_size
    shard_index    = rank % fsdp_group_size
    shard_ranks = list(range(shard_group_id * fsdp_group_size,
                             (shard_group_id + 1) * fsdp_group_size))
    fsdp_pg = dist.new_group(ranks=shard_ranks)

    replica_ranks = [shard_index + k * fsdp_group_size for k in range(num_replicas)]
    ddp_pg = dist.new_group(ranks=replica_ranks)

    ddp_rank_in_group = replica_ranks.index(rank)

    info = {
        "fsdp_pg": fsdp_pg,
        "ddp_pg": ddp_pg,
        "num_replicas": num_replicas,
        "ddp_rank": ddp_rank_in_group,
        "shard_group_id": shard_group_id,
        "shard_index": shard_index,
        "shard_ranks": shard_ranks,
        "replica_ranks": replica_ranks,
    }
    return info

def fsdp_wrap(module, sharding_strategy="full", mixed_precision=False, wrap_strategy="size", process_group=None, min_num_params=int(5e7), transformer_module=None, cpu_offload=False):
    if mixed_precision:
        mixed_precision_policy = MixedPrecision(
            param_dtype=torch.bfloat16,
            reduce_dtype=torch.float32,
            buffer_dtype=torch.float32,
            cast_forward_inputs=False
        )
    else:
        mixed_precision_policy = None

    if wrap_strategy == "transformer":
        auto_wrap_policy = partial(
            transformer_auto_wrap_policy,
            transformer_layer_cls=transformer_module
        )
    elif wrap_strategy == "size":
        auto_wrap_policy = partial(
            size_based_auto_wrap_policy,
            min_num_params=min_num_params
        )
    else:
        raise ValueError(f"Invalid wrap strategy: {wrap_strategy}")

    os.environ["NCCL_CROSS_NIC"] = "1"

    sharding_strategy = {
        "full": ShardingStrategy.FULL_SHARD,
        "hybrid_full": ShardingStrategy.HYBRID_SHARD,
        "hybrid_zero2": ShardingStrategy._HYBRID_SHARD_ZERO2,
        "no_shard": ShardingStrategy.NO_SHARD,
    }[sharding_strategy]

    module = FSDP(
        module,
        auto_wrap_policy=auto_wrap_policy,
        sharding_strategy=sharding_strategy,
        mixed_precision=mixed_precision_policy,
        device_id=torch.cuda.current_device(),
        limit_all_gathers=True,
        use_orig_params=True,
        cpu_offload=CPUOffload(offload_params=cpu_offload),
        sync_module_states=False,  # Load ckpt on rank 0 and sync to other ranks
        # process_group=process_group,
    )
    return module

def shard_model(
    model,
    *,
    cpu_offload: bool,
    reshard_after_forward: bool = True,
    mp_policy: MixedPrecisionPolicy | None = MixedPrecisionPolicy(),  # noqa
    mesh: DeviceMesh | None = None,
    fsdp_shard_conditions: list[Callable[[str, nn.Module], bool]] = [],  # noqa
    pin_cpu_memory: bool = True,
) -> None:
    """
    Utility to shard a model with FSDP using the PyTorch Distributed fully_shard API.
    This method will over the model's named modules from the bottom-up and apply shard modules
    based on whether they meet any of the criteria from shard_conditions.
    Args:
        model (TransformerDecoder): Model to shard with FSDP.
        shard_conditions (List[Callable[[str, nn.Module], bool]]): A list of functions to determine
            which modules to shard with FSDP. Each function should take module name (relative to root)
            and the module itself, returning True if FSDP should shard the module and False otherwise.
            If any of shard_conditions return True for a given module, it will be sharded by FSDP.
        cpu_offload (bool): If set to True, FSDP will offload parameters, gradients, and optimizer
            states to CPU.
        reshard_after_forward (bool): Whether to reshard parameters and buffers after
            the forward pass. Setting this to True corresponds to the FULL_SHARD sharding strategy
            from FSDP1, while setting it to False corresponds to the SHARD_GRAD_OP sharding strategy.
        mesh (Optional[DeviceMesh]): Device mesh to use for FSDP sharding under multiple parallelism.
            Default to None.
        fsdp_shard_conditions (List[Callable[[str, nn.Module], bool]]): A list of functions to determine
            which modules to shard with FSDP.
        pin_cpu_memory (bool): If set to True, FSDP will pin the CPU memory of the offloaded parameters.
    Raises:
        ValueError: If no layer modules were sharded, indicating that no shard_condition was triggered.
    """
    if fsdp_shard_conditions is None or len(fsdp_shard_conditions) == 0:
        logger.warning(
            "The FSDP shard condition list is empty or None. No modules will be sharded in %s",
            type(model).__name__)
        return

    fsdp_kwargs = {
        "reshard_after_forward": reshard_after_forward,
        "mesh": mesh,
        "mp_policy": mp_policy,
    }
    if cpu_offload:
        fsdp_kwargs["offload_policy"] = CPUOffloadPolicy(
            pin_memory=pin_cpu_memory)

    # iterating in reverse to start with
    # lowest-level modules first
    num_layers_sharded = 0
    # TODO(will): don't reshard after forward for the last layer to save on the
    # all-gather that will immediately happen Shard the model with FSDP,
    for n, m in reversed(list(model.named_modules())):
        if any([
                shard_condition(n, m)
                for shard_condition in fsdp_shard_conditions
        ]):
            fully_shard(m, **fsdp_kwargs)
            num_layers_sharded += 1

    if num_layers_sharded == 0:
        raise ValueError(
            "No layer modules were sharded. Please check if shard conditions are working as expected."
        )

    # Finally shard the entire model to account for any stragglers
    fully_shard(model, **fsdp_kwargs)


# TODO(PY): device mesh for cfg parallel
def load_model_from_full_model_state_dict(
    model: FSDPModule | torch.nn.Module,
    full_sd_iterator: Generator[tuple[str, torch.Tensor], None, None],
    device: torch.device,
    param_dtype: torch.dtype,
    strict: bool = False,
    cpu_offload: bool = False,
    param_names_mapping: Callable[[str], tuple[str, Any, Any]] | None = None,
    training_mode: bool = True,
) -> _IncompatibleKeys:
    """
    Converting full state dict into a sharded state dict
    and loading it into FSDP model (if training) or normal huggingface model
    Args:
        model (Union[FSDPModule, torch.nn.Module]): Model to generate fully qualified names for cpu_state_dict
        full_sd_iterator (Generator): an iterator yielding (param_name, tensor) pairs
        device (torch.device): device used to move full state dict tensors
        param_dtype (torch.dtype): dtype used to move full state dict tensors
        strict (bool): flag to check if to load the model in strict mode
        cpu_offload (bool): flag to check if FSDP offload is enabled
        param_names_mapping (Optional[Callable[[str], str]]): a function that maps full param name to sharded param name
        training_mode (bool): apply FSDP only for training
    Returns:
        ``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields:
            * **missing_keys** is a list of str containing the missing keys
            * **unexpected_keys** is a list of str containing the unexpected keys
    Raises:
        NotImplementedError: If got FSDP with more than 1D.
    """
    meta_sd = model.state_dict()
    sharded_sd = {}
    custom_param_sd, reverse_param_names_mapping = hf_to_custom_state_dict(
        full_sd_iterator, param_names_mapping)  # type: ignore
    for target_param_name, full_tensor in custom_param_sd.items():
        meta_sharded_param = meta_sd.get(target_param_name)
        if meta_sharded_param is None:
            raise ValueError(
                f"Parameter {target_param_name} not found in custom model state dict. The hf to custom mapping may be incorrect."
            )
        if not hasattr(meta_sharded_param, "device_mesh"):
            full_tensor = full_tensor.to(device=device, dtype=param_dtype)
            # In cases where parts of the model aren't sharded, some parameters will be plain tensors
            sharded_tensor = full_tensor
        else:
            full_tensor = full_tensor.to(device=device, dtype=param_dtype)
            sharded_tensor = distribute_tensor(
                full_tensor,
                meta_sharded_param.device_mesh,
                meta_sharded_param.placements,
            )
            if cpu_offload:
                sharded_tensor = sharded_tensor.cpu()
        sharded_sd[target_param_name] = nn.Parameter(sharded_tensor)

    model.reverse_param_names_mapping = reverse_param_names_mapping
    unused_keys = set(meta_sd.keys()) - set(sharded_sd.keys())
    if unused_keys:
        logger.warning("Found unloaded parameters in meta state dict: %s",
                       unused_keys)

    # List of allowed parameter name patterns
    ALLOWED_NEW_PARAM_PATTERNS = ["gate_compress"]  # Can be extended as needed
    for new_param_name in unused_keys:
        if not any(pattern in new_param_name
                   for pattern in ALLOWED_NEW_PARAM_PATTERNS):
            logger.error("Unsupported new parameter: %s. Allowed patterns: %s",
                         new_param_name, ALLOWED_NEW_PARAM_PATTERNS)
            raise ValueError(
                f"New parameter '{new_param_name}' is not supported. "
                f"Currently only parameters containing {ALLOWED_NEW_PARAM_PATTERNS} are allowed."
            )
        meta_sharded_param = meta_sd.get(new_param_name)
        if not hasattr(meta_sharded_param, "device_mesh"):
            # Initialize with zeros
            sharded_tensor = torch.zeros_like(meta_sharded_param,
                                              device=device,
                                              dtype=param_dtype)
        else:
            # Initialize with zeros and distribute
            full_tensor = torch.zeros_like(meta_sharded_param,
                                           device=device,
                                           dtype=param_dtype)
            sharded_tensor = distribute_tensor(
                full_tensor,
                meta_sharded_param.device_mesh,
                meta_sharded_param.placements,
            )
            if cpu_offload:
                sharded_tensor = sharded_tensor.cpu()
        sharded_sd[new_param_name] = nn.Parameter(sharded_tensor)

    # choose `assign=True` since we cannot call `copy_` on meta tensor
    return model.load_state_dict(sharded_sd, strict=strict, assign=True)

	# SPDX-License-Identifier: Apache-2.0

	# Adapted from torchtune
	# Copyright 2024 The TorchTune Authors.
	# Copyright 2025 The Trainer Authors.

	import contextlib
	from collections.abc import Callable, Generator
	from itertools import chain
	from typing import Any
	import os

	import torch
	import torch.distributed as dist
	from torch import nn
	from torch.distributed import DeviceMesh, init_device_mesh
	from torch.distributed._tensor import distribute_tensor
	from torch.distributed.fsdp import (CPUOffloadPolicy, FSDPModule,
	MixedPrecisionPolicy, fully_shard)
	from torch.nn.modules.module import _IncompatibleKeys
	from torch.nn.parallel import DistributedDataParallel as DDP

	from functools import partial
	from torch.distributed.fsdp import FullStateDictConfig, FullyShardedDataParallel as FSDP, MixedPrecision, ShardingStrategy, StateDictType
	from torch.distributed.fsdp.api import CPUOffload
	from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy, transformer_auto_wrap_policy

	from trainer.distributed import get_world_group
	from trainer.logger import init_logger
	from trainer.models.loader.utils import (get_param_names_mapping,
	hf_to_custom_state_dict)
	from trainer.models.loader.weight_utils import safetensors_weights_iterator
	from trainer.utils import set_mixed_precision_policy

	logger = init_logger(__name__)


	# TODO(PY): move this to utils elsewhere
	@contextlib.contextmanager
	def set_default_dtype(dtype: torch.dtype) -> Generator[None, None, None]:
	"""
	Context manager to set torch's default dtype.

	Args:
	dtype (torch.dtype): The desired default dtype inside the context manager.

	Returns:
	ContextManager: context manager for setting default dtype.

	Example:
	>>> with set_default_dtype(torch.bfloat16):
	>>> x = torch.tensor([1, 2, 3])
	>>> x.dtype
	torch.bfloat16


	"""
	old_dtype = torch.get_default_dtype()
	torch.set_default_dtype(dtype)
	try:
	yield
	finally:
	torch.set_default_dtype(old_dtype)


	# TODO(PY): add compile option
	def maybe_load_fsdp_model(
	load_from_dir: str,
	ar_action_load_from_dir: str,
	cls_name: str,
	model_cls: type[nn.Module],
	init_params: dict[str, Any],
	weight_dir_list: list[str],
	device: torch.device,
	hsdp_replicate_dim: int,
	hsdp_shard_dim: int,
	param_dtype: torch.dtype,
	reduce_dtype: torch.dtype,
	cpu_offload: bool = False,
	fsdp_inference: bool = False,
	output_dtype: torch.dtype \| None = None,
	training_mode: bool = True,
	pin_cpu_memory: bool = True,
	) -> torch.nn.Module:
	"""
	Load the model with FSDP if is training, else load the model without FSDP.
	"""
	# NOTE(will): cast_forward_inputs=True shouldn't be needed as we are
	# manually casting the inputs to the model
	mp_policy = MixedPrecisionPolicy(param_dtype,
	reduce_dtype,
	output_dtype,
	cast_forward_inputs=False)

	set_mixed_precision_policy(
	param_dtype=param_dtype,
	reduce_dtype=reduce_dtype,
	output_dtype=output_dtype,
	mp_policy=mp_policy,
	)

	# with set_default_dtype(param_dtype), torch.device("meta"):
	with set_default_dtype(param_dtype):
	logger.info(f"loading from: {load_from_dir}")
	logger.info(f"model:", model_cls)
	# model = model_cls(**init_params)
	model = model_cls.from_pretrained(
	load_from_dir, local_attn_size=-1, sink_size=0)


	if "Action" in cls_name:
	model.add_discrete_action_parameters()
	if ar_action_load_from_dir is not None:
	from safetensors.torch import load_file
	state_dict = load_file(ar_action_load_from_dir)
	model.load_state_dict(state_dict, strict=True)
	logger.info(f"loading from: {ar_action_load_from_dir}")
	# model.add_channel_concat_parameters()

	model.to(torch.float32)

	# Check if we should use FSDP
	use_fsdp = training_mode or fsdp_inference

	# Disable FSDP for MPS as it's not compatible
	from trainer.platforms import current_platform
	if current_platform.is_mps():
	use_fsdp = False
	logger.info("Disabling FSDP for MPS platform as it's not compatible")

	if use_fsdp:
	world_size = hsdp_replicate_dim * hsdp_shard_dim
	if not training_mode and not fsdp_inference:
	hsdp_replicate_dim = world_size
	hsdp_shard_dim = 1

	device_mesh = init_device_mesh(
	"cuda",
	# (Replicate(), Shard(dim=0))
	mesh_shape=(hsdp_replicate_dim, hsdp_shard_dim),
	mesh_dim_names=("replicate", "shard"),
	)
	shard_model(model,
	cpu_offload=cpu_offload,
	reshard_after_forward=True,
	mp_policy=mp_policy,
	mesh=device_mesh,
	fsdp_shard_conditions=model._fsdp_shard_conditions,
	pin_cpu_memory=pin_cpu_memory)

	for n, p in chain(model.named_parameters(), model.named_buffers()):
	if p.is_meta:
	raise RuntimeError(
	f"Unexpected param or buffer {n} on meta device.")
	# Avoid unintended computation graph accumulation during inference
	if isinstance(p, torch.nn.Parameter):
	p.requires_grad = False
	return model

	def build_2d_groups(rank: int, world_size: int, fsdp_group_size: int):
	assert world_size % fsdp_group_size == 0, "world_size 必须能整除 FSDP_GROUP_SIZE"
	num_replicas = world_size // fsdp_group_size

	shard_group_id = rank // fsdp_group_size
	shard_index = rank % fsdp_group_size
	shard_ranks = list(range(shard_group_id * fsdp_group_size,
	(shard_group_id + 1) * fsdp_group_size))
	fsdp_pg = dist.new_group(ranks=shard_ranks)

	replica_ranks = [shard_index + k * fsdp_group_size for k in range(num_replicas)]
	ddp_pg = dist.new_group(ranks=replica_ranks)

	ddp_rank_in_group = replica_ranks.index(rank)

	info = {
	"fsdp_pg": fsdp_pg,
	"ddp_pg": ddp_pg,
	"num_replicas": num_replicas,
	"ddp_rank": ddp_rank_in_group,
	"shard_group_id": shard_group_id,
	"shard_index": shard_index,
	"shard_ranks": shard_ranks,
	"replica_ranks": replica_ranks,
	}
	return info

	def fsdp_wrap(module, sharding_strategy="full", mixed_precision=False, wrap_strategy="size", process_group=None, min_num_params=int(5e7), transformer_module=None, cpu_offload=False):
	if mixed_precision:
	mixed_precision_policy = MixedPrecision(
	param_dtype=torch.bfloat16,
	reduce_dtype=torch.float32,
	buffer_dtype=torch.float32,
	cast_forward_inputs=False
	)
	else:
	mixed_precision_policy = None

	if wrap_strategy == "transformer":
	auto_wrap_policy = partial(
	transformer_auto_wrap_policy,
	transformer_layer_cls=transformer_module
	)
	elif wrap_strategy == "size":
	auto_wrap_policy = partial(
	size_based_auto_wrap_policy,
	min_num_params=min_num_params
	)
	else:
	raise ValueError(f"Invalid wrap strategy: {wrap_strategy}")

	os.environ["NCCL_CROSS_NIC"] = "1"

	sharding_strategy = {
	"full": ShardingStrategy.FULL_SHARD,
	"hybrid_full": ShardingStrategy.HYBRID_SHARD,
	"hybrid_zero2": ShardingStrategy._HYBRID_SHARD_ZERO2,
	"no_shard": ShardingStrategy.NO_SHARD,
	}[sharding_strategy]

	module = FSDP(
	module,
	auto_wrap_policy=auto_wrap_policy,
	sharding_strategy=sharding_strategy,
	mixed_precision=mixed_precision_policy,
	device_id=torch.cuda.current_device(),
	limit_all_gathers=True,
	use_orig_params=True,
	cpu_offload=CPUOffload(offload_params=cpu_offload),
	sync_module_states=False, # Load ckpt on rank 0 and sync to other ranks
	# process_group=process_group,
	)
	return module

	def shard_model(
	model,
	*,
	cpu_offload: bool,
	reshard_after_forward: bool = True,
	mp_policy: MixedPrecisionPolicy \| None = MixedPrecisionPolicy(), # noqa
	mesh: DeviceMesh \| None = None,
	fsdp_shard_conditions: list[Callable[[str, nn.Module], bool]] = [], # noqa
	pin_cpu_memory: bool = True,
	) -> None:
	"""
	Utility to shard a model with FSDP using the PyTorch Distributed fully_shard API.

	This method will over the model's named modules from the bottom-up and apply shard modules
	based on whether they meet any of the criteria from shard_conditions.

	Args:
	model (TransformerDecoder): Model to shard with FSDP.
	shard_conditions (List[Callable[[str, nn.Module], bool]]): A list of functions to determine
	which modules to shard with FSDP. Each function should take module name (relative to root)
	and the module itself, returning True if FSDP should shard the module and False otherwise.
	If any of shard_conditions return True for a given module, it will be sharded by FSDP.
	cpu_offload (bool): If set to True, FSDP will offload parameters, gradients, and optimizer
	states to CPU.
	reshard_after_forward (bool): Whether to reshard parameters and buffers after
	the forward pass. Setting this to True corresponds to the FULL_SHARD sharding strategy
	from FSDP1, while setting it to False corresponds to the SHARD_GRAD_OP sharding strategy.
	mesh (Optional[DeviceMesh]): Device mesh to use for FSDP sharding under multiple parallelism.
	Default to None.
	fsdp_shard_conditions (List[Callable[[str, nn.Module], bool]]): A list of functions to determine
	which modules to shard with FSDP.
	pin_cpu_memory (bool): If set to True, FSDP will pin the CPU memory of the offloaded parameters.

	Raises:
	ValueError: If no layer modules were sharded, indicating that no shard_condition was triggered.
	"""
	if fsdp_shard_conditions is None or len(fsdp_shard_conditions) == 0:
	logger.warning(
	"The FSDP shard condition list is empty or None. No modules will be sharded in %s",
	type(model).__name__)
	return

	fsdp_kwargs = {
	"reshard_after_forward": reshard_after_forward,
	"mesh": mesh,
	"mp_policy": mp_policy,
	}
	if cpu_offload:
	fsdp_kwargs["offload_policy"] = CPUOffloadPolicy(
	pin_memory=pin_cpu_memory)

	# iterating in reverse to start with
	# lowest-level modules first
	num_layers_sharded = 0
	# TODO(will): don't reshard after forward for the last layer to save on the
	# all-gather that will immediately happen Shard the model with FSDP,
	for n, m in reversed(list(model.named_modules())):
	if any([
	shard_condition(n, m)
	for shard_condition in fsdp_shard_conditions
	]):
	fully_shard(m, **fsdp_kwargs)
	num_layers_sharded += 1

	if num_layers_sharded == 0:
	raise ValueError(
	"No layer modules were sharded. Please check if shard conditions are working as expected."
	)

	# Finally shard the entire model to account for any stragglers
	fully_shard(model, **fsdp_kwargs)


	# TODO(PY): device mesh for cfg parallel
	def load_model_from_full_model_state_dict(
	model: FSDPModule \| torch.nn.Module,
	full_sd_iterator: Generator[tuple[str, torch.Tensor], None, None],
	device: torch.device,
	param_dtype: torch.dtype,
	strict: bool = False,
	cpu_offload: bool = False,
	param_names_mapping: Callable[[str], tuple[str, Any, Any]] \| None = None,
	training_mode: bool = True,
	) -> _IncompatibleKeys:
	"""
	Converting full state dict into a sharded state dict
	and loading it into FSDP model (if training) or normal huggingface model
	Args:
	model (Union[FSDPModule, torch.nn.Module]): Model to generate fully qualified names for cpu_state_dict
	full_sd_iterator (Generator): an iterator yielding (param_name, tensor) pairs
	device (torch.device): device used to move full state dict tensors
	param_dtype (torch.dtype): dtype used to move full state dict tensors
	strict (bool): flag to check if to load the model in strict mode
	cpu_offload (bool): flag to check if FSDP offload is enabled
	param_names_mapping (Optional[Callable[[str], str]]): a function that maps full param name to sharded param name
	training_mode (bool): apply FSDP only for training
	Returns:
	``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields:
	* missing_keys is a list of str containing the missing keys
	* unexpected_keys is a list of str containing the unexpected keys

	Raises:
	NotImplementedError: If got FSDP with more than 1D.
	"""
	meta_sd = model.state_dict()
	sharded_sd = {}
	custom_param_sd, reverse_param_names_mapping = hf_to_custom_state_dict(
	full_sd_iterator, param_names_mapping) # type: ignore
	for target_param_name, full_tensor in custom_param_sd.items():
	meta_sharded_param = meta_sd.get(target_param_name)
	if meta_sharded_param is None:
	raise ValueError(
	f"Parameter {target_param_name} not found in custom model state dict. The hf to custom mapping may be incorrect."
	)
	if not hasattr(meta_sharded_param, "device_mesh"):
	full_tensor = full_tensor.to(device=device, dtype=param_dtype)
	# In cases where parts of the model aren't sharded, some parameters will be plain tensors
	sharded_tensor = full_tensor
	else:
	full_tensor = full_tensor.to(device=device, dtype=param_dtype)
	sharded_tensor = distribute_tensor(
	full_tensor,
	meta_sharded_param.device_mesh,
	meta_sharded_param.placements,
	)
	if cpu_offload:
	sharded_tensor = sharded_tensor.cpu()
	sharded_sd[target_param_name] = nn.Parameter(sharded_tensor)

	model.reverse_param_names_mapping = reverse_param_names_mapping
	unused_keys = set(meta_sd.keys()) - set(sharded_sd.keys())
	if unused_keys:
	logger.warning("Found unloaded parameters in meta state dict: %s",
	unused_keys)

	# List of allowed parameter name patterns
	ALLOWED_NEW_PARAM_PATTERNS = ["gate_compress"] # Can be extended as needed
	for new_param_name in unused_keys:
	if not any(pattern in new_param_name
	for pattern in ALLOWED_NEW_PARAM_PATTERNS):
	logger.error("Unsupported new parameter: %s. Allowed patterns: %s",
	new_param_name, ALLOWED_NEW_PARAM_PATTERNS)
	raise ValueError(
	f"New parameter '{new_param_name}' is not supported. "
	f"Currently only parameters containing {ALLOWED_NEW_PARAM_PATTERNS} are allowed."
	)
	meta_sharded_param = meta_sd.get(new_param_name)
	if not hasattr(meta_sharded_param, "device_mesh"):
	# Initialize with zeros
	sharded_tensor = torch.zeros_like(meta_sharded_param,
	device=device,
	dtype=param_dtype)
	else:
	# Initialize with zeros and distribute
	full_tensor = torch.zeros_like(meta_sharded_param,
	device=device,
	dtype=param_dtype)
	sharded_tensor = distribute_tensor(
	full_tensor,
	meta_sharded_param.device_mesh,
	meta_sharded_param.placements,
	)
	if cpu_offload:
	sharded_tensor = sharded_tensor.cpu()
	sharded_sd[new_param_name] = nn.Parameter(sharded_tensor)

	# choose `assign=True` since we cannot call `copy_` on meta tensor
	return model.load_state_dict(sharded_sd, strict=strict, assign=True)