Neuron distributed training validation

finetune_overfit.py

@@ -0,0 +1,221 @@
+import os
+import shutil
+import tempfile
+from dataclasses import dataclass
+
+import torch
+from accelerate import ParallelismConfig
+from datasets import load_dataset
+
+from transformers import (
+    AutoConfig,
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    DataCollatorForLanguageModeling,
+    HfArgumentParser,
+    Trainer,
+    TrainingArguments,
+)
+from transformers.distributed import DistributedConfig
+from transformers.distributed.utils import _distributed_barrier, _ensure_torch_distributed
+from transformers.utils import is_torch_neuron_available
+
+
+@dataclass
+class ScriptArguments:
+    dataset_name: str
+    dataset_config: str | None = None
+    dataset_split: str = "train"
+    num_examples: int = 16  # tiny fixed subset to overfit on
+    max_length: int = 1024
+
+
+@dataclass
+class ModelArguments:
+    model_name_or_path: str
+    model_revision: str = "main"
+    trust_remote_code: bool = False
+    use_lora: bool = False
+
+
+def main(script_args, training_args, model_args):
+    if not torch.cuda.is_available() and is_torch_neuron_available(check_device=True):
+        import torch_neuronx  # noqa: F401
+
+    tp_size = int(os.environ.get("TP_SIZE", "1"))
+    fsdp_size = int(os.environ.get("FSDP_SIZE", "1"))
+    if tp_size > 1 and fsdp_size > 1:
+        raise ValueError(
+            f"TP_SIZE ({tp_size}) > 1 together with FSDP_SIZE ({fsdp_size}) > 1 is not supported: "
+            "1D `DistributedConfig` only. Use one or the other."
+        )
+    if (tp_size > 1 or fsdp_size > 1) and training_args.fsdp:
+        raise ValueError(
+            f"TP_SIZE ({tp_size}) / FSDP_SIZE ({fsdp_size}) together with --fsdp is not supported: "
+            "that flag configures Accelerate's own FSDP plugin, a separate mechanism from "
+            "`distributed_config`. Use one or the other."
+        )
+
+    kwargs = {}
+    if tp_size > 1:
+        training_args.parallelism_config = ParallelismConfig(tp_size=tp_size)
+        kwargs["distributed_config"] = DistributedConfig(tp_size=tp_size)
+    elif fsdp_size > 1:
+        kwargs["distributed_config"] = DistributedConfig(fsdp_size=fsdp_size)
+
+    if tp_size > 1 or fsdp_size > 1:
+        _ensure_torch_distributed()
+
+    config = AutoConfig.from_pretrained(
+        model_args.model_name_or_path,
+        revision=model_args.model_revision,
+        trust_remote_code=model_args.trust_remote_code,
+    )
+    dtype = torch.bfloat16 if training_args.bf16 else torch.float32
+
+    from peft import LoraConfig, get_peft_model, PeftModel
+
+    lora_config = LoraConfig(
+        r=64,
+        lora_alpha=128,
+        target_modules=[
+            "q_proj",
+            "k_proj",
+            "v_proj",
+            "o_proj",
+            "gate_proj",
+            "up_proj",
+            "down_proj",
+            "embed_tokens",
+        ],
+    )
+
+    # Loading the model once to make a copy of the original state dict for later comparison.
+    # No sharding is performed.
+    if model_args.use_lora:
+        original_model_path = os.path.join(training_args.output_dir, "original_model")
+        if int(os.environ.get("RANK", "0")) == 0:
+            original_model = AutoModelForCausalLM.from_pretrained(
+                model_args.model_name_or_path,
+                config=config,
+                revision=model_args.model_revision,
+                trust_remote_code=model_args.trust_remote_code,
+                torch_dtype=dtype,
+            )
+
+            original_model = get_peft_model(original_model, lora_config)
+
+            shutil.rmtree(original_model_path, ignore_errors=True)
+            original_model.save_pretrained(original_model_path)
+
+        if tp_size > 1 or fsdp_size > 1:
+            _distributed_barrier()
+
+        model = AutoModelForCausalLM.from_pretrained(
+            model_args.model_name_or_path,
+            config=config,
+            revision=model_args.model_revision,
+            trust_remote_code=model_args.trust_remote_code,
+            torch_dtype=dtype,
+            low_cpu_mem_usage=True,
+            **kwargs,
+        )
+
+        torch_device = f"cuda:{os.environ.get('LOCAL_RANK')}" if torch.cuda.is_available() else None
+        model = PeftModel.from_pretrained(
+            model, original_model_path, is_trainable=True, torch_device=torch_device
+        )
+        if int(os.environ.get("RANK", "0")) == 0:
+            model.print_trainable_parameters()
+    else:
+        model = AutoModelForCausalLM.from_pretrained(
+            model_args.model_name_or_path,
+            config=config,
+            revision=model_args.model_revision,
+            trust_remote_code=model_args.trust_remote_code,
+            torch_dtype=dtype,
+            low_cpu_mem_usage=True,
+            **kwargs,
+        )
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_args.model_name_or_path,
+        revision=model_args.model_revision,
+        trust_remote_code=model_args.trust_remote_code,
+    )
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    # ---------------------------------------------------------------------------
+    # Dataset: a tiny fixed subset, tokenized once via the chat template.
+    # ---------------------------------------------------------------------------
+    dataset = load_dataset(script_args.dataset_name, name=script_args.dataset_config, split=script_args.dataset_split)
+    dataset = dataset.select(range(script_args.num_examples))
+
+    def tokenize(example):
+        input_ids = tokenizer.apply_chat_template(
+            example["messages"],
+            tokenize=True,
+            return_dict=False,
+            truncation=True,
+            max_length=script_args.max_length,
+        )
+        return {"input_ids": input_ids}
+
+    dataset = dataset.map(tokenize, remove_columns=dataset.column_names)
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=dataset,
+        data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
+    )
+
+    # Check that the model can be saved and reloaded correctly, even when sharded.
+    rank = int(os.environ.get("RANK", "0"))
+    trainer.save_model(training_args.output_dir)
+
+    if rank == 0:
+        if model_args.use_lora:
+            original_state_dict = original_model.state_dict()
+
+            unsharded_model = PeftModel.from_pretrained(
+                AutoModelForCausalLM.from_pretrained(training_args.output_dir, torch_dtype=dtype),
+                original_model_path,
+            )
+            unsharded_state_dict = unsharded_model.state_dict()
+        else:
+            original_state_dict = original_model.state_dict()
+
+            unsharded_model = AutoModelForCausalLM.from_pretrained(training_args.output_dir, torch_dtype=dtype)
+            unsharded_state_dict = unsharded_model.state_dict()
+
+        mismatches = []
+        missing_keys = set(original_state_dict) - set(unsharded_state_dict)
+        if missing_keys:
+            mismatches.append(f"missing from saved checkpoint: {sorted(missing_keys)}")
+        for key, expected_value in unsharded_state_dict.items():
+            if key not in original_state_dict:
+                mismatches.append(f"{key}: unexpected key in saved checkpoint")
+                continue
+            try:
+                torch.testing.assert_close(expected_value, original_state_dict[key], rtol=0, atol=0)
+            except AssertionError as e:
+                mismatches.append(f"{key}: {e}")
+        if mismatches:
+            raise AssertionError("Save correctness check failed:\n" + "\n".join(mismatches))
+
+        print(
+            f"Save correctness check passed: {len(unsharded_state_dict)} parameters "
+            "match the unsharded checkpoint exactly."
+        )
+
+    # We can start training.
+    trainer.train()
+
+
+
+if __name__ == "__main__":
+    parser = HfArgumentParser((ScriptArguments, TrainingArguments, ModelArguments))
+    script_args, training_args, model_args = parser.parse_args_into_dataclasses()
+    main(script_args, training_args, model_args)

finetune_overfit.sh

@@ -0,0 +1,106 @@
+#!/bin/bash
+# Minimalist TP/FSDP overfitting test on AWS Trainium or CUDA GPUs using the regular
+# transformers Trainer (no SFTTrainer; LoRA optional via USE_LORA). Trains on a tiny fixed
+# subset of examples for many steps to check that the loss goes to ~0, as a correctness check
+# for TP or FSDP.
+
+set -euo pipefail
+
+# ---------------------------------------------------------------------------
+# Parallelism: pick a variant with PARALLEL_MODE=tp|fsdp (default: fsdp).
+# TP and FSDP2 are both enabled through the same mechanism -- `DistributedConfig`, which pre-shards
+# the model into DTensor at from_pretrained time:
+#   PARALLEL_MODE=tp   NUM_PROC=4 -> TP_SIZE=4
+#   PARALLEL_MODE=fsdp NUM_PROC=4 -> FSDP_SIZE=4
+# ---------------------------------------------------------------------------
+PARALLEL_MODE=${PARALLEL_MODE:-tp}
+NUM_PROC=${NUM_PROC:-2}
+USE_LORA=${USE_LORA:-true}
+
+case "$PARALLEL_MODE" in
+    tp)
+        export TP_SIZE=$NUM_PROC
+        export FSDP_SIZE=1
+        ;;
+    fsdp)
+        export TP_SIZE=1
+        export FSDP_SIZE=$NUM_PROC
+        ;;
+    *)
+        echo "Unknown PARALLEL_MODE: $PARALLEL_MODE (expected tp|fsdp)" >&2
+        exit 1
+        ;;
+esac
+
+# ---------------------------------------------------------------------------
+# Neuron runtime environment (Trainium only -- no-op, skipped entirely on CUDA)
+# ---------------------------------------------------------------------------
+if ! command -v nvidia-smi &> /dev/null || ! nvidia-smi &> /dev/null; then
+    export ON_NEURON_EAGER=1
+    export NEURON_EAGER_MODEL_CACHE_SIZE=10000
+    export OMP_NUM_THREADS=128
+    export HF_DEACTIVATE_ASYNC_LOAD=1
+
+    export TORCH_NEURONX_ENABLE_HOST_CC=1
+    export TORCH_NEURONX_ENABLE_ASYNC_NRT=1
+    #export NEURON_RT_NUM_CORES=1
+fi
+
+# ---------------------------------------------------------------------------
+# Model / data / hyperparameters
+# ---------------------------------------------------------------------------
+MODEL_NAME=Qwen/Qwen3-1.7B
+DATASET_NAME=trl-lib/Capybara
+LEARNING_RATE=5.0e-4
+NUM_TRAIN_EXAMPLES=16
+MAX_STEPS=50
+MAX_SEQ_LENGTH=1024
+BATCH_SIZE=4
+
+LORA_SUFFIX=""
+if [ "$USE_LORA" = "true" ]; then
+    LORA_SUFFIX="-lora"
+fi
+OUTPUT_DIR=Qwen3-1.7B-${PARALLEL_MODE}-Overfit${LORA_SUFFIX}
+
+
+echo "=========================================="
+echo "Plain Trainer parallelism overfitting test"
+echo "  Model:           $MODEL_NAME"
+echo "  Dataset:         $DATASET_NAME"
+echo "  PARALLEL_MODE:   $PARALLEL_MODE"
+echo "  NUM_PROC:        $NUM_PROC"
+echo "  TP_SIZE:         $TP_SIZE"
+echo "  FSDP_SIZE:       $FSDP_SIZE"
+echo "  USE_LORA:        $USE_LORA"
+echo "  Num examples:    $NUM_TRAIN_EXAMPLES"
+echo "  Max steps:       $MAX_STEPS"
+echo "  Batch:           $BATCH_SIZE"
+echo "  Max seq len:     $MAX_SEQ_LENGTH"
+echo "  Output dir:      $OUTPUT_DIR"
+echo "=========================================="
+
+if [ "$NUM_PROC" -eq 1 ]; then
+    LAUNCHER="python"
+else
+    LAUNCHER="torchrun --nproc_per_node=${NUM_PROC}"
+fi
+
+$LAUNCHER \
+    finetune_overfit.py \
+    --model_name_or_path "$MODEL_NAME" \
+    --use_lora $USE_LORA \
+    --dataset_name "$DATASET_NAME" \
+    --num_examples $NUM_TRAIN_EXAMPLES \
+    --max_length $MAX_SEQ_LENGTH \
+    --learning_rate $LEARNING_RATE \
+    --gradient_checkpointing true \
+    --bf16 true \
+    --per_device_train_batch_size $BATCH_SIZE \
+    --max_steps $MAX_STEPS \
+    --eval_strategy no \
+    --logging_steps 10 \
+    --save_strategy no \
+    --dataloader_num_workers 0 \
+    --report_to trackio \
+    --output_dir "$OUTPUT_DIR"

src/transformers/distributed/mixin.py

@@ -196,6 +196,7 @@ def maybe_distribute_model(
                 if isinstance(distributed_config.tp_plan, dict):
                     model.tp_plan = distributed_config.tp_plan
                 model = apply_tensor_parallelism(model, tp_mesh)
+                model._tp_size = distributed_config.tp_size
 
             elif distributed_config.fsdp_size > 1:
                 fsdp_mesh = device_mesh["fsdp"] if device_mesh.ndim > 1 else device_mesh