repository·main·Indexed 21 days ago
https://github.com/yandex/yafsdpA Sharded Data Parallelism framework optimized for transformer-based architectures like LLMs. YaFSDP aims to reduce communication and memory overhead, offering up to 20% faster pre-training compared to standard FSDP. It integrates with the Hugging Face stack, providing support for Causal Language Modeling (CLM) pre-training and Supervised Fine-Tuning (SFT) via accelerate launch. Key features include the fully_shard function, MixedPrecisionPolicy for optimized memory usage, and support for ragged tensors through RaggedShard and RaggedShardDTensor.
YaFSDP is a Sharded Data Parallelism framework specifically designed for transformer-like neural network architectures. It is optimized to reduce communication and memory operation overhead, making it particularly effective for Large Language Model (LLM) training.
Key advantages over standard FSDP include:
YaFSDP provides ready-to-use examples for training LLMs using the Hugging Face (🤗) stack. These examples are located in the examples directory:
examples/clm.md for instructions on causal language modeling pre-training.examples/sft.md for instructions on supervised fine-tuning workflows.The provided LLM training examples require a specific Docker image environment. This image is based on the NVIDIA PyTorch image and includes necessary patches for Hugging Face (🤗) libraries.
You can build the required image using the docker/build.sh script located in the repository.
./docker/build.shWhen migrating from standard PyTorch FSDP to YaFSDP, the primary interface change involves how modules are wrapped for sharding. YaFSDP replaces the auto_wrap_policy with a more explicit configuration involving module names, layer norm identification, and layer norm types. Additionally, you must explicitly provide the number of gradient_accumulation_steps to YaFSDP.
auto_wrap_policy, use:modules_to_wrap_with_names: A list of tuples (module, name) specifying which modules to shard. The names provided here are used in the state dict.rogue_layer_norm_modules_with_names: A dictionary mapping the first layer after transformer blocks (which typically contains only layer norm parameters) to its name.layer_norm_module_cls: The specific class type of the layer norm layers used in your model.sharding_strategy is mapped to a zero_stage integer (e.g., 3 for FULL_SHARD, 2 for SHARD_GRAD_OP).gradient_accumulation_steps as an argument.model: LlamaForCausalLM = ...
YaFSDP(
model,
zero_stage={
"ShardingStrategy.FULL_SHARD": 3,
"ShardingStrategy.SHARD_GRAD_OP": 2
}[sharding_strategy],
modules_to_wrap_with_names=[
(model.model.embed_tokens, "model.embed_tokens"),
*((m, f"model.layers.{i}") for i, m in enumerate(model.model.layers)),
(model.lm_head, "lm_head")
],
rogue_layer_norm_modules_with_names={model.norm: "model.norm"},
layer_norm_module_cls=LlamaRMSNorm,
param_dtype=param_dtype,
sync_module_states=sync_module_states,
param_init_fn=param_init_fn,
device_id=device,
gradient_accumulation_steps=gradient_accumulation_steps,
)The fsdp_config.yaml file used by accelerate launch controls low-level FSDP/YaFSDP behaviors. Key configuration keys include:
| Key | Description |
|---|---|
fsdp_state_dict_type | Choose between FULL_STATE_DICT (global gathered state) or LOCAL_STATE_DICT (local sharded states). |
fsdp_activation_checkpointing | Boolean toggle for activation checkpointing. |
fsdp_num_layers_to_checkpoint | Integer specifying the number of layers to checkpoint. |
num_processes | Total number of training processes (calculated as number of hosts x number of devices per host). |
YaFSDP automates the unshard/reshard lifecycle through PyTorch hooks:
_pre_forward hook: Triggers root_pre_forward (handling input casting and stream synchronization) and YaFSDPParamGroup.pre_forward (unsharding parameters)._post_forward hook: Triggers YaFSDPParamGroup.post_forward (resharding/cleanup) and registers a pre-backward hook on the output tensors._pre_backward hook: Triggered by the autograd engine on tensors requiring gradients. It manages pre_backward logic for parameter groups and handles prefetching for the next iteration._root_post_backward_final_callback: A callback queued via the execution engine that runs once per backward pass to finalize states, clean up communication contexts, and handle the is_last_backward logic.YaFSDP uses a custom torch.autograd.Function called RegisterPostBackwardFunction to automate post-backward operations.
When _register_post_backward_hook is called on a YaFSDPParamGroup, it identifies tensors in the function arguments that require gradients and wraps them with this autograd function. During the backward pass, once the gradients for these tensors are computed, the backward method of RegisterPostBackwardFunction is triggered, which calls param_group.post_backward(). This allows YaFSDP to perform necessary cleanup or communication (like reduce-scatter) immediately after the gradients are available.
class RegisterPostBackwardFunction(torch.autograd.Function):
@staticmethod
def forward(
ctx: Any, param_group: YaFSDPParamGroup, *inputs: torch.Tensor
) -> tuple[torch.Tensor, ...]:
ctx.param_group = param_group
return inputs
@staticmethod
def backward(
ctx: Any, *grads: torch.Tensor
) -> tuple[None, *tuple[torch.Tensor, ...]]:
ctx.param_group.post_backward()
return (None, *grads)YaFSDP uses a hierarchical state management system to coordinate distributed training.
YaFSDPState (subclass of torch.distributed._composable_state._State): Represents the state for a specific module or group of modules. It manages parameter groups (YaFSDPParamGroup), communication contexts (YaFSDPCommContext), and training lifecycle states (e.g., FORWARD, PRE_BACKWARD, POST_BACKWARD).YaFSDPStateContext: A shared context object used by YaFSDPState instances to coordinate across the module tree. It tracks all states in the root tree (all_states), identifies the current iteration's forward root (iter_forward_root), and manages synchronization events like post_optim_event to ensure all-gather streams wait for the optimizer to complete.YaFSDP supports user-defined metadata that is passed through the all-gather process via ExtensionsData. This is useful for handling complex operations like FP8 quantization scales that need to be synchronized during the all-gather step.
If the underlying tensor has fsdp_pre_all_gather and fsdp_post_all_gather methods defined, YaFSDPParam will automatically initialize an ExtensionsData object to hold this metadata.
Usage Pattern:
set_all_gather_input(): Triggers the fsdp_pre_all_gather hook to prepare metadata.init_unsharded_param(...): Triggers the fsdp_post_all_gather hook to apply metadata to the newly unsharded tensor.The TrainingState enum tracks the lifecycle of a training step, which is critical for managing transitions between sharding, unsharding, and gradient reduction.
States:
FORWARD: Covers the period from the start of the forward pass (pre-forward) until the end of the forward pass (post-forward).PRE_BACKWARD: Occurs when unsharding happens during the backward pass.POST_BACKWARD: Occurs when resharding and reducing gradients happens after the backward pass.IDLE: The state before/after a forward pass, or before PRE_BACKWARD/after POST_BACKWARD.from ya_fsdp._common import TrainingState
# Example usage in a training loop logic
current_state = TrainingState.IDLE
# ... transition to FORWARD ...If use_yccl=True is passed to YaFSDP, the library uses YCCL for high-performance 16-bit non-layer-norm all-gathers and reduce-scatters.
When using YCCL, the library automatically handles padding to ensure compatibility. If you need to manually specify the intra- and inter-color mapping for YCCL (e.g., for custom network topologies), use the yccl_intra_inter_expr argument, which accepts a string containing a Python expression that evaluates to (intra_color, inter_color) using the current rank.
You can perform distributed supervised fine-tuning by launching a Docker container equipped with ya-fsdp. This setup integrates with the Hugging Face ecosystem, specifically using trl, transformers, and accelerate.
To enable YaFSDP within an accelerate launch command, you must:
--config_file (e.g., ya-fsdp/examples/fsdp_config.yaml).--fsdp_ya_fsdp_enabled true to activate the YaFSDP backend.This example uses the trl/examples/scripts/sft.py script to fine-tune a Llama-3 model on the openassistant-guanaco dataset.
docker run \
-it \
--rm \
--net host \
--gpus '"device=0,1"' \
--ipc=host \
--ulimit memlock=-1 \
--ulimit stack=67108864 \
ya-fsdp:latest \
accelerate launch \
--config_file ya-fsdp/examples/fsdp_config.yaml \
--fsdp_ya_fsdp_enabled true \
trl/examples/scripts/sft.py \
--do_train \
--model_name_or_path meta-llama/Meta-Llama-3-8B \
--max_steps 5 \
--block_size 2048 \
--per_device_train_batch_size 1 \
--per_device_eval_batch_size 1 \
--dataset_name timdettmers/openassistant-guanaco \
--save_strategy no \
--logging_steps 1 \
--report_to tensorboard \
--output_dir sft