examples/pytorch/3d_parallel_checks.py

# Copyright 2024 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
""":
This script is used to test training a model using Tensor Parallelism and Data Parallelism.

Usage:
export CUDA_VISIBLE_DEVICES=0,1,2,3
export CUDA_VISIBLE_DEVICES=4,5,6,7
export CUDA_VISIBLE_DEVICES=5,6,7
TP_SIZE=2 DP_SIZE=2 torchrun --nproc_per_node=4 --rdzv_endpoint=localhost:29503 test_train.py
CP_SIZE=2 DP_SIZE=2 torchrun --nproc_per_node=4 test_train.py
CP_SIZE=2 TP_SIZE=2 torchrun --nproc_per_node=4 test_train.py

TP_SIZE=1 CP_SIZE=4 torchrun --nproc_per_node=4 test_train.py
TP_SIZE=1 DP_SIZE=4 torchrun --nproc_per_node=4 test_train.py
TP_SIZE=4 DP_SIZE=1 torchrun --nproc_per_node=4 --rdzv_endpoint=localhost:29503 test_train.py
IGNORE_SANITY=1 CP_SIZE=1 TP_SIZE=1 DP_SIZE=1 torchrun --nproc_per_node=1 --rdzv_endpoint=l
ocalhost:29504 test_train.py
"""

import logging
import os
from collections.abc import Iterable
from contextlib import nullcontext

import torch
import torch.distributed as dist
import torch.distributed.checkpoint as dcp
import torch.nn as nn
import torch.optim as optim
import wandb
from datasets import load_dataset
from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict
from torch.distributed.checkpoint.stateful import Stateful
from torch.distributed.device_mesh import DeviceMesh
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import ShardingStrategy
from torch.distributed.tensor import DTensor
from torch.distributed.tensor.experimental import context_parallel
from torch.nn.attention import SDPBackend, sdpa_kernel
from torch.utils.data import DataLoader, default_collate
from torch.utils.data.distributed import DistributedSampler

from transformers import AutoModelForCausalLM, AutoTokenizer


ignore_sanity_checks = int(os.environ.get("IGNORE_SANITY", "0")) == 1
# torch.use_deterministic_algorithms(True)
torch.backends.cudnn.deterministic = True

# Set up logging
logging.basicConfig(
    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
    datefmt="%m/%d/%Y %H:%M:%S",
    level=logging.INFO,
)
logger = logging.getLogger(__name__)

# from torch.distributed.tensor.experimental._attention import set_rotate_method

# set_rotate_method("alltoall")  # rotate shards using all-to-all


def main():
    tp_size = int(os.environ.get("TP_SIZE", "1"))
    dp_size = int(os.environ.get("DP_SIZE", "4"))
    cp_size = int(os.environ.get("CP_SIZE", "1"))  # Add CP size configuration
    sdpa_backend = SDPBackend.FLASH_ATTENTION  # For CP
    # sdpa_backend = SDPBackend.MATH # For CP
    global_batch_size = 8  # Desired global batch size
    seq_len = 1024  # Sequence length
    num_train_steps = 10000  # Number of training steps
    LR = 1e-5
    model_name = "HuggingFaceTB/SmolLM2-1.7B"
    # model_name = "unsloth/Llama-3.2-1B"

    CHECKPOINT_DIR = f"checkpoint_tp{tp_size}_dp{dp_size}_cp{cp_size}"

    # Initialize distributed environment
    if "RANK" in os.environ and "WORLD_SIZE" in os.environ:
        dist.init_process_group("nccl")
        rank = dist.get_rank()
        world_size = dist.get_world_size()
        local_rank = int(os.environ["LOCAL_RANK"])
        torch.cuda.set_device(local_rank)

        assert world_size == tp_size * dp_size * cp_size, (
            f"World size ({world_size}) must equal TP size ({tp_size}) * DP size ({dp_size}) * CP size ({cp_size})"
        )

        mesh = torch.arange(world_size).reshape(dp_size, tp_size, cp_size)
        world_mesh = DeviceMesh(device_type="cuda", mesh=mesh, mesh_dim_names=("dp", "tp", "cp"))
        tp_mesh = world_mesh["tp"]
        dp_mesh = world_mesh["dp"]
        cp_mesh = world_mesh["cp"]
        world_mesh["dp", "cp"]._flatten(mesh_dim_name="dp_cp")
        logger.info(f"Created DeviceMesh: {world_mesh}")
        logger.info(
            f"Distributed setup - Rank: {rank}, World size: {world_size}, Local rank: {local_rank}, DP: {dp_mesh.get_local_rank()}, TP: {tp_mesh.get_local_rank()}, CP: {cp_mesh.get_local_rank()}"
        )

        if dist.get_rank() == 0:
            wandb.init(
                project="tp_dp_test",
                config={
                    "tp_size": tp_size,
                    "dp_size": dp_size,
                    "cp_size": cp_size,
                    "global_batch_size": global_batch_size,
                    "model_name": model_name,
                    "dataset": "roneneldan/TinyStories-1M",
                    "seq_len": seq_len,
                    "lr": LR,
                    "weight_decay": 0.1,
                },
                name=f"llama_tp{tp_size}_dp{dp_size}_cp{cp_size}"
                if model_name == "unsloth/Llama-3.2-1B"
                else f"tp{tp_size}_dp{dp_size}_cp{cp_size}",
            )
            logger.info(f"ignore_sanity_checks is set to: {ignore_sanity_checks}")
            logger.info("Wandb initialized.")
            # Log the current file to wandb
            wandb.save("test_train.py")

    else:
        logger.info("Running in non-distributed mode. DeviceMesh not applicable.")
        rank = 0
        world_size = 1
        local_rank = 0
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        wandb.init(
            project="tp_dp_test",
            config={
                "tp_size": 1,
                "dp_size": 1,
                "global_batch_size": global_batch_size,
                "model_name": model_name,
                "dataset": "roneneldan/TinyStories-1M",
                "seq_len": seq_len,
            },
            name="llama_tp1_dp1_nondist" if model_name == "unsloth/Llama-3.2-1B" else "tp1_dp1_nondist",
        )
        logger.info("Wandb initialized for non-distributed run.")

    # Load model and tokenizer
    logger.info(f"Loading model and tokenizer from {model_name}")
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
        logger.info(f"Set pad_token to eos_token: {tokenizer.pad_token}")

    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        device_mesh=tp_mesh if dist.is_initialized() else None,
        tp_plan="auto",
        dtype=torch.bfloat16,
    )
    logger.info(f"Model loaded onto device mesh: {tp_mesh}")

    if dist.is_initialized():
        assert model.config.num_key_value_heads % tp_mesh.size() == 0, (
            f"num_key_value_heads={model.config.num_key_value_heads} must be divisible by tp_size={tp_mesh.size()}"
        )
        device = torch.device(f"cuda:{local_rank}")
    else:
        model = model.to(device)

    logger.info(f"Using device: {device} for non-model tensors")
    use_ddp = False
    if dist.is_initialized() and dp_mesh.size() > 1:
        # FSDP1
        model = FSDP(model, device_mesh=dp_mesh, sharding_strategy=ShardingStrategy.NO_SHARD)
        # FSDP2
        # for transformer_block in model.model.layers:
        #     fully_shard(transformer_block, mesh=dp_mesh, reshard_after_forward=False)
        # fully_shard(model.model, mesh=dp_mesh, reshard_after_forward=False)
        # DDP
        # replicate(model, device_mesh=dp_mesh, bucket_cap_mb=100)
        # assert len(list(model.parameters()))>5, "No parameters found in model. Probably DDP/FSDP bug.." # TODO: we should be cautious abt using model.parameters()
        use_ddp = True

    model.train()
    assert len(list(model.parameters())) > 0, "No parameters found in model. Probably DDP bug.."
    assert len([p for p in model.parameters() if p.requires_grad]) > 0, (
        "No gradients found in model. Probably DDP bug.."
    )

    if dist.is_initialized() and not ignore_sanity_checks:
        # assert model is replicated across all dp
        for name, param in model.named_parameters():
            sanity_check_tensor_sync(param, dp_mesh)

        # assert model is different across tp (only for sharded params)
        for name, param in model.named_parameters():
            if isinstance(param, DTensor) and param.placements[0].is_shard():
                # Only check sharded parameters for non-sync across TP
                sanity_check_tensor_sync(param, tp_mesh, not_sync=True)
            elif isinstance(param, DTensor) and param.placements[0].is_replicate():
                # Replicated parameters should be the same across TP
                sanity_check_tensor_sync(param, tp_mesh)

        # assert model is replicated across cp
        for name, param in model.named_parameters():
            sanity_check_tensor_sync(param, cp_mesh)

    # Load and preprocess TinyStories dataset
    logger.info("Loading TinyStories dataset...")
    raw_dataset = load_dataset("roneneldan/TinyStories", split="train[:1%]")  # Use 1% for faster testing

    def tokenize_function(examples):
        # Tokenize the text without padding
        tokenized_batch = tokenizer(
            examples["text"], padding=False, truncation=True, max_length=seq_len, return_tensors=None
        )
        # Set labels to be the same as input_ids for Causal LM
        tokenized_batch["labels"] = tokenized_batch["input_ids"].copy()
        return tokenized_batch

    tokenized_dataset = raw_dataset.map(tokenize_function, batched=True, remove_columns=["text"])
    logger.info(f"Dataset loaded and tokenized. Size: {len(tokenized_dataset)}")

    # Create packed sequences
    def create_packed_sequences(examples):
        # Flatten all sequences
        all_tokens = []
        for input_ids in examples["input_ids"]:
            all_tokens.extend(input_ids)

        # Split into sequences of seq_len + 1 (for input + label)
        num_sequences = len(all_tokens) // (seq_len + 1)
        packed_input_ids = []
        packed_labels = []

        for i in range(num_sequences):
            start_idx = i * (seq_len + 1)
            end_idx = start_idx + (seq_len + 1)
            # Get the full sequence
            full_sequence = all_tokens[start_idx:end_idx]
            # For input_ids, remove the last token
            packed_input_ids.append(full_sequence[:-1])
            # For labels, remove the first token
            packed_labels.append(full_sequence[1:])

        return {"input_ids": packed_input_ids, "labels": packed_labels}

    # Apply packing to the dataset
    packed_dataset = tokenized_dataset.map(
        create_packed_sequences,
        batched=True,
        remove_columns=tokenized_dataset.column_names,
        batch_size=1000,  # Process in batches for efficiency
        num_proc=60,
    )
    logger.info(f"Dataset packed. New size: {len(packed_dataset)}")

    # Shuffle the packed dataset
    packed_dataset = packed_dataset.shuffle(seed=42)
    logger.info("Packed dataset shuffled")

    # Calculate local batch size
    if dist.is_initialized():
        assert global_batch_size % dp_mesh.size() == 0, (
            f"Global batch size ({global_batch_size}) must be divisible by DP size ({dp_mesh.size()})"
        )
        local_batch_size = global_batch_size // dp_mesh.size()
    else:
        local_batch_size = global_batch_size

    logger.info(
        f"Global batch size: {global_batch_size}, DP size: {dp_size if dist.is_initialized() else 1}, Local batch size: {local_batch_size}"
    )

    # Simple collate function since sequences are already packed
    def collate_fn(batch):
        input_ids = torch.tensor([item["input_ids"] for item in batch], dtype=torch.long)
        labels = torch.tensor([item["labels"] for item in batch], dtype=torch.long)
        return {"input_ids": input_ids, "labels": labels}

    if dist.is_initialized():
        sampler = DistributedSampler(
            packed_dataset, num_replicas=dp_mesh.size(), rank=dp_mesh.get_local_rank(), shuffle=False
        )
    else:
        sampler = None

    dataloader = DataLoader(
        packed_dataset,
        batch_size=local_batch_size,
        sampler=sampler,
        shuffle=False,
        collate_fn=collate_fn,
    )
    logger.info(f"DataLoader created. Distributed: {dist.is_initialized()}")

    optimizer = optim.AdamW(model.parameters(), lr=LR, weight_decay=0.1)

    # Training loop
    logger.info(f"Starting training for {num_train_steps} steps...")
    model.train()
    step = 0
    while step < num_train_steps:
        for batch in dataloader:
            if step >= num_train_steps:
                break  # Exit loop if max steps reached

            # Move batch to appropriate device
            batch = {k: v.to(device) for k, v in batch.items()}

            # Sanity checks for batch distribution (only if distributed)
            if dist.is_initialized() and not ignore_sanity_checks:
                # check batch is same across all tp
                sanity_check_tensor_sync(batch["input_ids"], tp_mesh)
                # check batch is different across dp
                sanity_check_tensor_sync(batch["input_ids"], dp_mesh, not_sync=True)

            optimizer.zero_grad()

            # Add position_ids to batch before CP sharding
            batch_size = batch["input_ids"].shape[0]
            position_ids = torch.arange(0, seq_len, dtype=torch.long, device=device)
            position_ids = position_ids.unsqueeze(0).expand(batch_size, -1)
            batch["position_ids"] = position_ids
            from torch.distributed.tensor.experimental._attention import _cp_options

            _cp_options.enable_load_balance = False

            with sdpa_kernel(sdpa_backend):  # TODO: ideally move this to attention implementation
                cp_context = (
                    nullcontext()
                    if cp_mesh.size() == 1
                    else context_parallel(
                        cp_mesh,
                        buffers=[
                            batch["input_ids"],
                            batch["labels"],
                            batch["position_ids"],
                        ],  # TODO: need to add attention mask
                        buffer_seq_dims=[1, 1, 1],
                    )
                )
                with cp_context:
                    # Pop labels from batch before model forward pass
                    labels = batch.pop("labels")
                    outputs = model(**batch)  # [mbs, seq_len/cp]
                    loss = outputs.loss
                    logits = outputs.logits

                    # Compute loss with shifted labels
                    loss = model.loss_function(
                        logits=logits, labels=None, shift_labels=labels, vocab_size=model.config.vocab_size
                    )

                    # Sanity checks for logits
                    if dist.is_initialized() and not ignore_sanity_checks:
                        # sanity_check_tensor_sync(logits, tp_mesh) # TODO: only true without sequence parallel
                        sanity_check_tensor_sync(logits, dp_mesh, not_sync=True)
                        sanity_check_tensor_sync(logits, cp_mesh, not_sync=True)

                    loss.backward()

                # all reduce grads across dp_cp if applicable
                all_reduce_grads(model, world_mesh, use_ddp=use_ddp)

                # Sanity checks for gradients (only if distributed)
                if dist.is_initialized() and not ignore_sanity_checks:
                    # check grads are not same across all tp (for sharded grads)
                    for name, param in model.named_parameters():
                        if param.grad is not None and isinstance(param.grad, DTensor):
                            if param.grad.placements[0].is_shard():
                                sanity_check_tensor_sync(param.grad, tp_mesh, not_sync=True)
                            elif param.grad.placements[0].is_replicate():
                                sanity_check_tensor_sync(param.grad, tp_mesh)
                    # check grads are same across dp
                    for name, param in model.named_parameters():
                        if param.grad is not None and dp_mesh.size() > 1:
                            sanity_check_tensor_sync(param.grad, dp_mesh)
                    # check grads are same across cp
                    for name, param in model.named_parameters():
                        if param.grad is not None and cp_mesh.size() > 1:
                            sanity_check_tensor_sync(param.grad, cp_mesh)

                # Calculate gradient norm and clip gradients
                if hasattr(model, "clip_grad_norm_"):
                    # when using FSDP or DDP, model.parameters() doesn't work
                    gradnorm = model.clip_grad_norm_(max_norm=1.0, norm_type=2.0)
                else:
                    assert len(list(model.parameters())) > 2, "No parameters found in model. Probably DDP bug.."
                    assert len([p for p in model.parameters() if p.requires_grad]) > 2, (
                        "No gradients found in model. Probably DDP bug.."
                    )
                    assert len([p for p in model.parameters() if p.grad is not None]) > 2, (
                        "No gradients found in model. Probably DDP bug.."
                    )
                    # only works with FSDP's NO_SHARD otherwise we should use FSDP's clip_grad_norm_
                    gradnorm = clip_grad_norm_(model.parameters(), max_norm=1.0, norm_type=2.0, foreach=True)

                optimizer.step()
                # Sanity checks for updated model parameters (only if distributed)
                if dist.is_initialized() and not ignore_sanity_checks:
                    # check updated model is different across all tp (for sharded params)
                    for name, param in model.named_parameters():
                        if isinstance(param, DTensor):
                            if param.placements[0].is_shard():
                                sanity_check_tensor_sync(param, tp_mesh, not_sync=True)
                            elif param.placements[0].is_replicate():
                                sanity_check_tensor_sync(param, tp_mesh)
                    # check updated model is same across dp
                    for name, param in model.named_parameters():
                        sanity_check_tensor_sync(param, dp_mesh)
                    # check updated model is same across cp
                    for name, param in model.named_parameters():
                        sanity_check_tensor_sync(param, cp_mesh)

                # allreduce loss across cp_dp before logging
                if dist.is_initialized() and (cp_mesh.size() > 1 or dp_mesh.size() > 1):
                    dist.all_reduce(loss, group=world_mesh["dp_cp"].get_group(), op=dist.ReduceOp.AVG)
                current_loss = loss.item()

                # Log loss and gradnorm to wandb (only on rank 0 of dp group)
                if not dist.is_initialized() or dist.get_rank() == 0:
                    logger.info(
                        f"Step: {step} | GBS: {global_batch_size} | DP: {dp_mesh.size()} | TP: {tp_mesh.size()} | CP: {cp_mesh.size()} | Loss: {current_loss} | Gradnorm: {gradnorm} | lr: {LR}"
                    )
                    wandb.log(
                        {
                            "train/loss": current_loss,
                            "train/gradnorm": gradnorm,
                            "step": step,
                            "lr": LR,
                            "GBS": global_batch_size,
                        }
                    )

            step += 1  # Increment step count

    logger.info("Training loop finished.")

    # Save model using DCP (only if distributed)
    if dist.is_initialized():
        state_dict = {"app": AppState(model, optimizer)}
        dcp.save(
            state_dict=state_dict,
            checkpoint_id=CHECKPOINT_DIR,
        )
        logger.info(f"Saved checkpoint to {CHECKPOINT_DIR}")
    else:
        # Fallback to regular save for non-distributed case
        save_dir = "test_model_nondist"
        model.save_pretrained(save_dir)
        tokenizer.save_pretrained(save_dir)  # Save tokenizer too
        logger.info(f"Saved model to {save_dir}")

    # Example of loading the checkpoint (only if distributed)
    if dist.is_initialized():
        # Create a new model instance
        logger.info("Creating new model instance for verification")
        new_model = AutoModelForCausalLM.from_pretrained(
            model_name,
            device_mesh=tp_mesh,
            dtype=torch.bfloat16,  # Use same dtype
        )
        new_optimizer = optim.AdamW(new_model.parameters(), lr=LR)

        # Load checkpoint into new model
        state_dict = {"app": AppState(new_model, new_optimizer)}
        dcp.load(
            state_dict=state_dict,
            checkpoint_id=CHECKPOINT_DIR,
        )
        logger.info("Loaded checkpoint into new model")

        # Verify model weights match
        logger.info("Verifying model weights match...")
        for (name1, param1), (name2, param2) in zip(model.named_parameters(), new_model.named_parameters()):
            torch.testing.assert_close(
                param1.to_local(),
                param2.to_local(),
                rtol=1e-3,
                atol=1e-3,
                msg=f"Weights mismatch in {name1} vs {name2}",
            )

        # Verify optimizer states match
        logger.info("Verifying optimizer states match...")
        for name1, state1 in optimizer.state_dict().items():
            state2 = new_optimizer.state_dict()[name1]
            if name1 == "state":
                # Compare state dictionaries for each parameter
                for param_id, param_state1 in state1.items():
                    param_state2 = state2[param_id]
                    # Compare each state component (step, exp_avg, exp_avg_sq)
                    for key, value1 in param_state1.items():
                        value2 = param_state2[key]
                        if isinstance(value1, DTensor):
                            # Convert DTensors to local tensors for comparison
                            torch.testing.assert_close(
                                value1.to_local(),
                                value2.to_local(),
                                rtol=1e-5,
                                atol=1e-5,
                                msg=f"Optimizer state mismatch in state[{param_id}][{key}]",
                            )
                        else:
                            torch.testing.assert_close(
                                value1,
                                value2,
                                rtol=1e-5,
                                atol=1e-5,
                                msg=f"Optimizer state mismatch in state[{param_id}][{key}]",
                            )
            elif name1 == "param_groups":
                # Compare param_groups (excluding the actual params list)
                for i, (group1, group2) in enumerate(zip(state1, state2)):
                    for key in group1:
                        if key != "params":  # Skip comparing the params list
                            assert group1[key] == group2[key], f"Param group mismatch in param_groups[{i}][{key}]"

        # Run a forward pass with both models to verify outputs match
        logger.info("Running forward pass verification...")
        with torch.no_grad():
            # Use the last batch for verification
            batch = {k: v.to(device) for k, v in batch.items()}  # Ensure batch is on correct device
            original_outputs = model(**batch)
            new_outputs = new_model(**batch)
            torch.testing.assert_close(
                original_outputs.logits.to_local(),
                new_outputs.logits.to_local(),
                rtol=1e-3,
                atol=1e-3,
                msg="Model outputs do not match!",
            )  # Increased tolerance slightly for bf16

    # Clean up distributed environment and finish wandb run
    if dist.is_initialized():
        dist.destroy_process_group()
        logger.info("Cleaned up distributed process group")
        # Finish wandb run on rank 0
        if dist.get_rank() == 0:
            wandb.finish()
            logger.info("Wandb run finished.")
    else:
        wandb.finish()
        logger.info("Wandb run finished.")


def all_reduce_grads(model, world_mesh, use_ddp):
    """All reduce gradients across dp_cp if applicable."""
    cp_mesh = world_mesh["cp"]
    if use_ddp:
        # DDP takes care of syncing grads
        mesh = cp_mesh
    else:
        mesh = world_mesh["dp", "cp"]._flatten(mesh_dim_name="dp_cp")
    if dist.is_initialized() and mesh.size() > 1:
        for name, param in model.named_parameters():
            if param.grad is not None:
                # Workaround for cross-mesh communication limitation with DTensor gradients
                if isinstance(param.grad, DTensor):
                    local_grad = param.grad.to_local()
                    # Ensure grad requires grad for inplace modification checks (might not be needed)
                    # local_grad = local_grad.detach().requires_grad_(True)
                    torch.distributed.all_reduce(local_grad, op=torch.distributed.ReduceOp.SUM, group=mesh.get_group())
                    local_grad = local_grad / mesh.size()
                    # Assign averaged grad back - need careful handling if DTensor structure is complex
                    # This simple assignment might work if the grad structure matches param structure
                    param.grad = DTensor.from_local(
                        local_grad, device_mesh=param.grad.device_mesh, placements=param.grad.placements
                    )
                else:
                    # Handle regular tensors if any exist (e.g. buffers not converted to DTensor)
                    torch.distributed.all_reduce(param.grad, op=torch.distributed.ReduceOp.AVG, group=mesh.get_group())


class ContextParallelCollator:
    """Collator for context parallel training that splits sequences into chunks."""

    def __init__(self, cp_mesh: DeviceMesh | None = None):
        self.cp_mesh = cp_mesh

    def __call__(self, batch: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
        batch = default_collate(batch)
        if self.cp_mesh is not None and self.cp_mesh.size() > 1:
            # Get sequence length from the input batch
            seq_len = batch["input_ids"].shape[1]
            assert seq_len % self.cp_mesh.size() == 0, (
                f"Sequence length {seq_len} must be divisible by CP size {self.cp_mesh.size()}"
            )
            chunk_size = seq_len // self.cp_mesh.size()
            cp_rank = self.cp_mesh.get_local_rank()
            start_idx = cp_rank * chunk_size
            end_idx = start_idx + chunk_size

            # Keep only the local chunk of the sequence
            batch["input_ids"] = batch["input_ids"][:, start_idx:end_idx]
            batch["attention_mask"] = batch["attention_mask"][:, start_idx:end_idx]
            batch["labels"] = batch["labels"][:, start_idx:end_idx]

        return batch


class AppState(Stateful):
    """Wrapper for checkpointing the Application State including model and optimizer."""

    def __init__(self, model, optimizer=None):
        self.model = model
        self.optimizer = optimizer

    def state_dict(self):
        model_state_dict, optimizer_state_dict = get_state_dict(self.model, self.optimizer)
        return {"model": model_state_dict, "optim": optimizer_state_dict}

    def load_state_dict(self, state_dict):
        set_state_dict(
            self.model, self.optimizer, model_state_dict=state_dict["model"], optim_state_dict=state_dict["optim"]
        )


def sanity_check_tensor_sync(
    tensor: torch.Tensor, mesh: DeviceMesh, rtol: float = 1e-4, atol: float = 1e-4, not_sync: bool = False
) -> None:
    """
    Verify that a tensor is synchronized (or not synchronized) across all processes in the mesh's process group.
    Handles both regular tensors and DTensors.

    Args:
        tensor (torch.Tensor): The tensor to check for synchronization (can be DTensor)
        mesh (DeviceMesh): The device mesh containing the process group
        rtol (float): Relative tolerance for comparison
        atol (float): Absolute tolerance for comparison
        not_sync (bool): If True, asserts that tensors are NOT synchronized. If False, asserts they are synchronized.
    """
    if not dist.is_initialized() or mesh.size() == 1:
        return  # No need to check in non-distributed mode

    # Get the process group from the mesh
    pg = mesh.get_group()

    # Convert DTensor to local tensor if needed
    if hasattr(tensor, "to_local"):
        local_tensor = tensor.to_local()
    else:
        local_tensor = tensor

    # Gather tensors from all processes
    world_size = dist.get_world_size(pg)
    gathered_tensors = [torch.empty_like(local_tensor) for _ in range(world_size)]
    dist.all_gather(gathered_tensors, local_tensor, group=pg)

    # Compare each tensor with the first one
    for i in range(1, world_size):
        try:
            torch.testing.assert_close(gathered_tensors[0], gathered_tensors[i], rtol=rtol, atol=atol)
        except AssertionError as e:
            if not_sync:
                continue
            # # Add detailed debugging for logit synchronization issues
            # print(f"\nLogit synchronization error between rank 0 and rank {i}:")
            # print(f"Tensor shape: {gathered_tensors[0].shape}")
            # print(f"Number of mismatched elements: {(gathered_tensors[0] != gathered_tensors[i]).sum()}")
            # print(f"Percentage of mismatched elements: {((gathered_tensors[0] != gathered_tensors[i]).sum() / gathered_tensors[0].numel() * 100):.2f}%")

            # # Find the first few mismatches
            # mismatches = torch.nonzero(gathered_tensors[0] != gathered_tensors[i])
            # print("\nFirst few mismatches:")
            # for idx in mismatches[:5]:
            #     idx = tuple(idx.tolist())
            #     print(f"Index {idx}:")
            #     print(f"Rank 0 value: {gathered_tensors[0][idx]}")
            #     print(f"Rank {i} value: {gathered_tensors[i][idx]}")
            #     print(f"Absolute difference: {abs(gathered_tensors[0][idx] - gathered_tensors[i][idx])}")
            #     print(f"Relative difference: {abs(gathered_tensors[0][idx] - gathered_tensors[i][idx]) / max(abs(gathered_tensors[0][idx]), abs(gathered_tensors[i][idx]))}")

            # # Check if differences are systematic (e.g., all positive or negative)
            # diff = gathered_tensors[0] - gathered_tensors[i]
            # print(f"\nDifference statistics:")
            # print(f"Mean difference: {diff.mean()}")
            # print(f"Std difference: {diff.std()}")
            # print(f"Max positive difference: {diff.max()}")
            # print(f"Max negative difference: {diff.min()}")
            raise e


def clip_grad_norm_(
    parameters: Iterable[torch.Tensor],
    max_norm: float,
    norm_type: float = 2.0,
    error_if_nonfinite: bool = False,
    foreach: bool | None = None,
) -> torch.Tensor:
    """
    Clip the gradient norm of an iterable of parameters.
    """
    # Filter out parameters with no gradients
    parameters = [p for p in parameters if p.grad is not None]
    assert len(parameters) > 0, "No parameters with gradients found"

    # Calculate total norm
    if norm_type == float("inf"):
        total_norm = max(p.grad.detach().abs().max() for p in parameters)
    else:
        total_norm = torch.norm(torch.stack([torch.norm(p.grad.detach(), norm_type) for p in parameters]), norm_type)

    # Convert DTensor to local tensor if needed
    if isinstance(total_norm, DTensor):
        total_norm = total_norm.full_tensor()

    # Clip gradients
    clip_coef = max_norm / (total_norm + 1e-6)
    if clip_coef < 1:
        for p in parameters:
            p.grad.detach().mul_(clip_coef)

    return total_norm


def check_params_sync(model_params, original_params):
    """
    Check if original_params are being updated in sync with model parameters.

    Args:
        model_params: Iterator of model parameters after update
        original_params: List of original parameters before DDP wrapping
    """
    for mp, op in zip(model_params, original_params):
        if isinstance(mp, DTensor):
            mp = mp.to_local()
        if isinstance(op, DTensor):
            op = op.to_local()
        if not torch.allclose(mp.data, op.data, rtol=0, atol=0):
            raise RuntimeError(f"Parameters out of sync: model param {mp.data} != original param {op.data}")
    return True


def get_parameters(model: nn.Module) -> Iterable[torch.Tensor]:
    """
    Get all parameters from a model by iterating over its modules.
    This is an alternative to model.parameters() that works with DTensor models.

    Args:
        model (nn.Module): The model to get parameters from

    Returns:
        Iterable[torch.Tensor]: An iterator over all parameters in the model
    """
    for module in model._modules.values():
        # Look for parameters in module attributes
        for attr in module.__dict__.values():
            if isinstance(attr, torch.Tensor) and attr.requires_grad:
                yield attr
        # Recursively get parameters from submodules
        yield from get_parameters(module)


def update_model_parameters(model: nn.Module) -> None:
    """
    Update model._parameters using named_modules() to ensure all parameters are properly tracked.

    Args:
        model (nn.Module): The model to update parameters for
    """
    # Clear existing parameters
    model._parameters = {}

    # Add parameters from named_modules
    for name, module in model.named_modules():
        # Skip the root module itself
        if name == "":
            continue

        # Get the parameter name by removing 'module.' prefix if it exists
        param_name = name.replace("module.", "")

        # Add weight and bias parameters if they exist
        if hasattr(module, "weight") and module.weight is not None:
            model._parameters[f"{param_name}.weight"] = module.weight
        if hasattr(module, "bias") and module.bias is not None:
            model._parameters[f"{param_name}.bias"] = module.bias


if __name__ == "__main__":
    main()