### Launch Distributed SFT with YAFSDP (Bash)

Source: https://github.com/yandex/yafsdp/blob/main/examples/sft.md

This command launches a distributed supervised fine-tuning job using YAFSDP, accelerating the process with libraries like trl, transformers, and accelerate. It requires Docker and specifies GPU devices, configuration files, and training parameters for the model and dataset.

```bash
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
```

--------------------------------

### Sharding Model with PyTorch FSDP

Source: https://github.com/yandex/yafsdp/blob/main/docs/migration_guide.md

Demonstrates how to shard a LlamaForCausalLM model using PyTorch's FSDP. It configures sharding strategy, auto-wrapping policy for specific layers, mixed precision, and prefetching. This setup requires PyTorch and its distributed training components.

```python
model: LlamaForCausalLM = ...

FSDP(
    model,
    sharding_strategy=sharding_strategy,
    auto_wrap_policy=partial(
        lambda_auto_wrap_policy,
        lambda_fn=lambda m: (
            m is model.model.embed_tokens
            or m in model.model.layers
            or m is model.lm_head
        ),
    ),
    mixed_precision=MixedPrecision(param_dtype=param_dtype, reduce_dtype=torch.float32),
    sync_module_states=sync_module_states,
    param_init_fn=param_init_fn,
    device_id=device,
    backward_prefetch=BackwardPrefetch.BACKWARD_PRE,
    forward_prefetch=True,
    use_orig_params=True,
)
```

--------------------------------

### Launch Distributed Causal LM Pre-training with Docker and YaFSDP

Source: https://github.com/yandex/yafsdp/blob/main/examples/clm.md

This command initiates a distributed Causal Language Model pre-training job. It utilizes Docker for environment setup and `accelerate` for distributed training, with YaFSDP enabled for efficient model parallelism. Key parameters control GPU usage, model and tokenizer selection, training steps, batch sizes, dataset, and output directory.

```bash
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 \
        transformers/examples/pytorch/language-modeling/run_clm.py \
            --do_train \
            --config_name meta-llama/Meta-Llama-3-8B \
            --tokenizer_name 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 wikitext \
            --dataset_config_name wikitext-2-raw-v1 \
            --save_strategy no \
            --logging_steps 1 \
            --report_to tensorboard \
            --output_dir clm
```

--------------------------------

### Sharding Model with YaFSDP

Source: https://github.com/yandex/yafsdp/blob/main/docs/migration_guide.md

Shows an equivalent model sharding configuration using YaFSDP for a LlamaForCausalLM model. It replaces `auto_wrap_policy` with explicit module lists and layer norm configurations, and requires specifying gradient accumulation steps. Dependencies include YaFSDP and PyTorch.

```python
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,
)
```

--------------------------------

### Initialize YaFSDP Wrapper for Distributed Training

Source: https://context7.com/yandex/yafsdp/llms.txt

Initializes the YaFSDP wrapper to shard model parameters across distributed processes. It supports ZeRO-3 sharding, mixed precision with bfloat16, gradient accumulation, and specific module wrapping. Inputs include the model and distributed environment setup; outputs are a wrapped model ready for training.

```python
import torch
import torch.nn as nn
from ya_fsdp import YaFSDP
from transformers import LlamaForCausalLM, LlamaConfig

# Initialize model
config = LlamaConfig.from_pretrained("meta-llama/Meta-Llama-3-8B")
model = LlamaForCausalLM(config)

# Create process groups for distributed training
import torch.distributed as dist
dist.init_process_group(backend="nccl")

# Wrap model with YaFSDP
wrapped_model = YaFSDP(
    module=model,
    zero_stage=3,  # ZeRO-3 sharding (use 2 for ZeRO-2)
    param_dtype=torch.bfloat16,
    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")
    ],
    layer_norm_module_cls=type(model.model.layers[0].input_layernorm),
    gradient_accumulation_steps=4,
    data_parallel_process_group=dist.group.WORLD,
    bit16_reduce_scatter=True,
    bit32_acc_for_bit16_reduce_scatter=True,
    hpz_first_layers_num=0,
    sync_module_states=True,
    device_id=torch.cuda.current_device()
)

# Forward pass
input_ids = torch.randint(0, 32000, (2, 512)).cuda()
output = wrapped_model(input_ids=input_ids)

# Backward pass
loss = output.logits.sum()
loss.backward()

# Gradient clipping
total_norm = wrapped_model.clip_grad_norm_(max_norm=1.0)
print(f"Gradient norm: {total_norm}")
```

--------------------------------

### Enable YCCL for optimized all-to-all collectives with YaFSDP

Source: https://context7.com/yandex/yafsdp/llms.txt

This code snippet shows how to enable YCCL for optimized all-to-all collectives within the YaFSDP framework. It requires YCCL installation and allows for pipelined all-gather operations, bit16 reduce-scatter, and profiling. Note that using YCCL automatically disables HPZ.

```python
from yafs.torch.fsdp import YaFSDP
import torch

# Assuming 'model', 'intra_node_group', 'data_parallel_group', 'model_parallel_group', and 'LayerNormClass' are defined elsewhere
model_with_yccl = YaFSDP(
    model,
    zero_stage=3,
    param_dtype=torch.bfloat16,
    modules_to_wrap_with_names=[...],
    layer_norm_module_cls=LayerNormClass,
    gradient_accumulation_steps=4,
    use_yccl=True,
    yccl_all_gather_stages=2,  # Pipeline stages for all-gather
    bit16_reduce_scatter=True,
    bit32_acc_for_bit16_reduce_scatter=True,
    profile=True,  # Enable YCCL profiling
    data_parallel_process_group=data_parallel_group,
    intra_node_data_parallel_process_group=intra_node_group,
    model_parallel_process_group=model_parallel_group
)
```

--------------------------------

### Launch Supervised Fine-Tuning with YaFSDP (Bash)

Source: https://context7.com/yandex/yafsdp/llms.txt

This bash script initiates a supervised fine-tuning process for language models using YaFSDP and HuggingFace TRL. It launches a Docker container and runs the `accelerate launch` command with parameters tailored for instruction-based fine-tuning.

```bash
# Launch supervised fine-tuning
docker run \
    -it --rm --net host \
    --gpus all \
    --ipc=host \
    --ulimit memlock=-1 \
    --ulimit stack=67108864 \
    -v /path/to/data:/data \
    ya-fsdp:latest \
    accelerate launch \
        --config_file ya-fsdp/examples/fsdp_config.yaml \
        --fsdp_ya_fsdp_enabled true \
        --num_processes 16 \
        trl/examples/scripts/sft.py \
            --do_train \
            --model_name_or_path meta-llama/Meta-Llama-3-70B \
            --max_steps 5000 \
            --max_seq_length 2048 \
            --per_device_train_batch_size 1 \
            --gradient_accumulation_steps 8 \
            --learning_rate 1e-5 \
            --dataset_name timdettmers/openassistant-guanaco \
            --save_strategy steps \
            --save_steps 500 \
            --logging_steps 5 \
            --bf16 true \
            --gradient_checkpointing true \
            --output_dir /data/sft_output
```

--------------------------------

### Launch Causal Language Model Pre-Training with YaFSDP (Bash)

Source: https://context7.com/yandex/yafsdp/llms.txt

This bash script launches a distributed training job for causal language model pre-training using YaFSDP. It involves building a Docker image, running the container, and executing the `accelerate launch` command with specific configurations for HuggingFace Transformers.

```bash
# Build Docker image with YaFSDP
cd docker && bash build.sh

# Launch distributed training
docker run \
    -it --rm --net host \
    --gpus '"device=0,1,2,3,4,5,6,7"' \
    --ipc=host \
    --ulimit memlock=-1 \
    --ulimit stack=67108864 \
    -v /path/to/data:/data \
    ya-fsdp:latest \
    accelerate launch \
        --config_file ya-fsdp/examples/fsdp_config.yaml \
        --fsdp_ya_fsdp_enabled true \
        --num_processes 8 \
        transformers/examples/pytorch/language-modeling/run_clm.py \
            --do_train \
            --model_name_or_path meta-llama/Meta-Llama-3-8B \
            --max_steps 10000 \
            --block_size 4096 \
            --per_device_train_batch_size 1 \
            --gradient_accumulation_steps 16 \
            --learning_rate 3e-4 \
            --dataset_name wikitext \
            --dataset_config_name wikitext-103-raw-v1 \
            --save_strategy steps \
            --save_steps 1000 \
            --logging_steps 10 \
            --bf16 true \
            --output_dir /data/clm_output
```

--------------------------------

### Initialize Process Groups for HPZ and YCCL Optimization (Python)

Source: https://context7.com/yandex/yafsdp/llms.txt

This Python snippet shows the initialization of distributed process groups required for Hierarchical Parameter Sharding (HPZ) and Yandex Collective Communications Library (YCCL) optimizations. It demonstrates creating data parallel and intra-node groups, with an option for model parallel groups.

```python
from ya_fsdp import YaFSDP
import torch.distributed as dist

# Initialize process groups
data_parallel_group = dist.new_group()  # All ranks
intra_node_group = dist.new_group()     # Ranks on same node
model_parallel_group = None              # Optional: for tensor parallelism
```

--------------------------------

### Migration from PyTorch FSDP to YaFSDP (Python)

Source: https://context7.com/yandex/yafsdp/llms.txt

This Python code demonstrates how to migrate from PyTorch's native FSDP to YaFSDP. It shows the equivalent configuration for wrapping models, specifying sharding strategies, mixed precision, and other relevant parameters, highlighting the differences in API and options.

```python
# Original PyTorch FSDP code
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import MixedPrecision, ShardingStrategy
from torch.distributed.fsdp.wrap import lambda_auto_wrap_policy
from functools import partial

model = FSDP(
    model,
    sharding_strategy=ShardingStrategy.FULL_SHARD,
    auto_wrap_policy=partial(
        lambda_auto_wrap_policy,
        lambda_fn=lambda m: (
            m is model.model.embed_tokens
            or m in model.model.layers
            or m is model.lm_head
        ),
    ),
    mixed_precision=MixedPrecision(
        param_dtype=torch.bfloat16,
        reduce_dtype=torch.float32
    ),
    sync_module_states=True,
    device_id=torch.cuda.current_device(),
    use_orig_params=True
)

# Equivalent YaFSDP code
from ya_fsdp import YaFSDP

wrapped_model = YaFSDP(
    model,
    zero_stage=3,  # FULL_SHARD = ZeRO-3, SHARD_GRAD_OP = ZeRO-2
    param_dtype=torch.bfloat16,
    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.model.norm: "model.norm"
    },
    layer_norm_module_cls=type(model.model.layers[0].input_layernorm),
    gradient_accumulation_steps=4,
    sync_module_states=True,
    device_id=torch.cuda.current_device(),
    bit16_reduce_scatter=True,
    bit32_acc_for_bit16_reduce_scatter=True
)
```

--------------------------------

### Enable HPZ for first N layers with YaFSDP

Source: https://context7.com/yandex/yafsdp/llms.txt

This code snippet demonstrates how to enable HPZ (intra-node all-gather optimization) for the first N layers of a model using YaFSDP. It configures the FSDP wrapper with specific parameters like zero_stage, param_dtype, and gradient_accumulation_steps, while enabling HPZ for a specified number of layers.

```python
from yafs.torch.fsdp import YaFSDP
import torch

# Assuming 'model', 'intra_node_group', 'data_parallel_group', and 'LayerNormClass' are defined elsewhere
model = YaFSDP(
    model,
    zero_stage=3,
    param_dtype=torch.bfloat16,
    modules_to_wrap_with_names=[...],
    layer_norm_module_cls=LayerNormClass,
    gradient_accumulation_steps=4,
    hpz_first_layers_num=8,  # Apply HPZ to first 8 transformer layers
    intra_node_data_parallel_process_group=intra_node_group,
    data_parallel_process_group=data_parallel_group,
    device_id=torch.cuda.current_device()
)
```

--------------------------------

### Manage State Dicts with YaFSDP

Source: https://context7.com/yandex/yafsdp/llms.txt

Save and load model checkpoints using YaFSDP, supporting both full and sharded state dictionaries. This allows for flexible checkpointing strategies, including gathering the full state dict on rank 0 or saving individual shards per rank, with options for offloading to CPU.

```python
from ya_fsdp import fully_shard, StateDictType, FullStateDictConfig, ShardedStateDictConfig
import torch
import torch.distributed as dist

# Initialize model with YaFSDP
model = fully_shard(model, mesh=mesh, mp_policy=mp_policy)

# Save full state dict (gathered on rank 0)
with model.state_dict_type(
    StateDictType.FULL_STATE_DICT,
    FullStateDictConfig(rank0_only=True, offload_to_cpu=True)
):
    state_dict = model.state_dict()
    if dist.get_rank() == 0:
        torch.save(state_dict, "checkpoint_full.pt")

# Save sharded state dict (each rank saves its shard)
with model.state_dict_type(
    StateDictType.SHARDED_STATE_DICT,
    ShardedStateDictConfig(offload_to_cpu=True)
):
    state_dict = model.state_dict()
    rank = dist.get_rank()
    torch.save(state_dict, f"checkpoint_shard_rank{rank}.pt")

# Load state dict
checkpoint = torch.load("checkpoint_full.pt", map_location="cpu")
model.load_state_dict(checkpoint)

# Context manager for temporary state dict type change
with model.state_dict_type(StateDictType.FULL_STATE_DICT):
    # Operations using full state dict
    state = model.state_dict()
# Automatically reverts to previous state dict type
```

--------------------------------

### Configure YaFSDP MixedPrecisionPolicy

Source: https://context7.com/yandex/yafsdp/llms.txt

Configures mixed precision training by defining parameter, reduction, and output data types. This policy allows for fine-grained control over precision to balance performance and numerical stability. It can be applied using the `fully_shard` function for end-to-end FP16 training or BF16 with FP32 reduction.

```python
from ya_fsdp import MixedPrecisionPolicy
import torch

# BF16 parameters with FP32 gradient reduction
policy = MixedPrecisionPolicy(
    param_dtype=torch.bfloat16,
    reduce_dtype=torch.float32,
    output_dtype=None,  # Keep output in compute dtype
    cast_forward_inputs=False,
    bit32_acc_for_bit16_reduce_scatter=True
)

# Apply to model with fully_shard
from ya_fsdp import fully_shard
from torch.distributed.device_mesh import init_device_mesh

mesh = init_device_mesh("cuda", (torch.distributed.get_world_size(),))
model = fully_shard(
    model, 
    mesh=mesh,
    mp_policy=policy,
    reshard_after_forward=True
)

# FP16 end-to-end training
fp16_policy = MixedPrecisionPolicy(
    param_dtype=torch.float16,
    reduce_dtype=torch.float16,
    cast_forward_inputs=True
)
```

--------------------------------

### Set Post-Optimizer Event for Precise Synchronization (PyTorch)

Source: https://context7.com/yandex/yafsdp/llms.txt

This code snippet demonstrates how to set a post-optimizer event for precise synchronization in PyTorch training loops using `torch.cuda.Event`. It's crucial for accurate performance profiling and synchronization.

```python
import torch.cuda

post_optim_event = torch.cuda.Event()

for batch in dataloader:
    outputs = model(input_ids=batch["input_ids"].cuda())
    loss = outputs.loss
    loss.backward()

    optimizer.step()
    post_optim_event.record()
    model.set_post_optim_event(post_optim_event)

    optimizer.zero_grad()
    post_optim_event = torch.cuda.Event()  # Create new event for next iteration
```

--------------------------------

### Configure Advanced Prefetching with YaFSDP

Source: https://context7.com/yandex/yafsdp/llms.txt

Configure explicit prefetching in YaFSDP to overlap communication with computation during forward and backward passes. This involves setting specific modules to prefetch in advance, optimizing performance by reducing idle time.

```python
from ya_fsdp import fully_shard
import torch.nn as nn

# Assume model has sequential layers
layers = list(model.model.layers)

# Apply YaFSDP to each layer
for layer in layers:
    fully_shard(layer, mesh=mesh, mp_policy=mp_policy)
fully_shard(model, mesh=mesh, mp_policy=mp_policy)

# Configure forward prefetching: each layer prefetches next 2 layers
for i, layer in enumerate(layers[:-2]):
    layer.set_modules_to_forward_prefetch([layers[i+1], layers[i+2]])

# Last layers only prefetch next layer
if len(layers) > 1:
    layers[-2].set_modules_to_forward_prefetch([layers[-1]])

# Configure backward prefetching: reverse order
for i in range(2, len(layers)):
    layer = layers[i]
    layer.set_modules_to_backward_prefetch([layers[i-1], layers[i-2]])
```

--------------------------------

### Manually Unshard and Reshard Parameters

Source: https://context7.com/yandex/yafsdp/llms.txt

Explicitly control parameter sharding in YaFSDP, useful for inference or evaluation modes where unsharded parameters are preferred. This allows for manual un-sharding before inference and re-sharding afterwards to manage memory efficiently.

```python
from ya_fsdp import fully_shard
import torch

model = fully_shard(model, mesh=mesh)

# Evaluation mode: keep parameters unsharded
model.eval()
model.set_reshard_after_forward(reshard_after_forward=False, recurse=True)

with torch.no_grad():
    # Manually unshard all parameters
    model.unshard()

    # Run inference with unsharded parameters
    for batch in eval_dataloader:
        outputs = model(input_ids=batch["input_ids"].cuda())
        # Process outputs...

    # Optionally reshard to free memory
    model.reshard()

# Training mode: reshard after forward
model.train()
model.set_reshard_after_forward(reshard_after_forward=True, recurse=True)

# Control backward resharding for gradient accumulation
model.set_reshard_after_backward(
    reshard_after_backward=False,  # Keep unsharded during accumulation
    recurse=True
)
```

--------------------------------

### Control Gradient Accumulation Synchronization

Source: https://context7.com/yandex/yafsdp/llms.txt

Configure gradient synchronization for efficient gradient accumulation in YaFSDP. This prevents communication overhead during intermediate accumulation steps by disabling gradient synchronization, enabling it only for the final step.

```python
from ya_fsdp import fully_shard
import torch

model = fully_shard(model, mesh=mesh, mp_policy=mp_policy)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

gradient_accumulation_steps = 4

for step, batch in enumerate(dataloader):
    is_accumulation_step = (step + 1) % gradient_accumulation_steps != 0

    # Disable gradient sync on accumulation steps
    model.set_requires_gradient_sync(
        requires_gradient_sync=not is_accumulation_step,
        recurse=True
    )

    # Forward and backward
    input_ids = batch["input_ids"].cuda()
    outputs = model(input_ids=input_ids, labels=batch["labels"].cuda())
    loss = outputs.loss / gradient_accumulation_steps
    loss.backward()

    # Update weights after accumulation completes
    if not is_accumulation_step:
        optimizer.step()
        optimizer.zero_grad()

# Set last backward flag for cleanup
model.set_is_last_backward(is_last_backward=True)
```

--------------------------------

### Apply YaFSDP using PyTorch 2.x fully_shard Function

Source: https://context7.com/yandex/yafsdp/llms.txt

Applies YaFSDP to modules using PyTorch 2.x composable distributed API, enabling automatic state management. This method is suitable for integrating with PyTorch 2.x features like device meshes and mixed precision policies. It shards individual layers, embeddings, and the model head, followed by sharding the root module.

```python
import torch
import torch.nn as nn
from torch.distributed.device_mesh import init_device_mesh
from ya_fsdp import fully_shard, MixedPrecisionPolicy
from transformers import LlamaForCausalLM

# Initialize distributed environment
torch.distributed.init_process_group(backend="nccl")
mesh = init_device_mesh("cuda", (torch.distributed.get_world_size(),))

# Create model
model = LlamaForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B",
    torch_dtype=torch.bfloat16
)

# Apply fully_shard to each layer
mp_policy = MixedPrecisionPolicy(
    param_dtype=torch.bfloat16,
    reduce_dtype=torch.float32
)

for layer in model.model.layers:
    fully_shard(
        layer, 
        mesh=mesh,
        reshard_after_forward=True,
        mp_policy=mp_policy
    )

# Shard embeddings and head
fully_shard(model.model.embed_tokens, mesh=mesh, mp_policy=mp_policy)
fully_shard(model.lm_head, mesh=mesh, mp_policy=mp_policy)

# Apply to root module
fully_shard(model, mesh=mesh, mp_policy=mp_policy)

# Training loop
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
model.train()

for batch in dataloader:
    input_ids = batch["input_ids"].cuda()
    labels = batch["labels"].cuda()

    outputs = model(input_ids=input_ids, labels=labels)
    loss = outputs.loss

    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.