diff --git a/training/deepspeed_finetune_demo/README.md b/training/deepspeed_finetune_demo/README.md index d39d1bee6..d5a167e9b 100644 --- a/training/deepspeed_finetune_demo/README.md +++ b/training/deepspeed_finetune_demo/README.md @@ -113,6 +113,49 @@ With AutoEP, each rank holds a different expert shard. The training script saves Use `convert_ds_to_hf.py` to merge all shards back into a standard HuggingFace model. +## Validate AutoEP checkpoint resume + +`run_autoep_affine_ir_checkpoint_experiment.sh` compares the loss from uninterrupted +training with loss after resuming from both a native DeepSpeed checkpoint and a +Universal checkpoint. It runs each path through step 100, saves a native +checkpoint at step 50, converts that checkpoint to Universal format, and compares +the losses for steps 51-100. + +Run the experiment on a Linux host with the required GPUs and dependencies +installed. Set `DEEPSPEED_REPO` to the root of a DeepSpeed checkout that provides +`deepspeed/checkpoint/ds_to_universal.py`. The model and dataset must already be +cached because Transformers and Datasets offline mode are enabled by default: + +```bash +cd training/deepspeed_finetune_demo +export DEEPSPEED_REPO=/path/to/DeepSpeed +./run_autoep_affine_ir_checkpoint_experiment.sh +``` + +The script defaults to 8 GPUs, AutoEP size 8, ZeRO stage 2, the +`moonshotai/Moonlight-16B-A3B` model, and the `tatsu-lab/alpaca` dataset. Override +these settings with environment variables as needed: + +```bash +NUM_GPUS=8 AUTOEP_SIZE=8 ZERO_STAGE=2 \ +MODEL_NAME=moonshotai/Moonlight-16B-A3B \ +DATASET_NAME=tatsu-lab/alpaca \ +OUTPUT_ROOT=/path/to/experiment \ +./run_autoep_affine_ir_checkpoint_experiment.sh +``` + +`DEEPSPEED_LAUNCHER` defaults to `ds`; set it to `deepspeed` if that is the +launcher available in your environment. To use GPUs selected by DeepSpeed's +`--include` option, set `GPU_INCLUDE` (for example, `GPU_INCLUDE=localhost:0,1`). +To allow downloads instead of using cached model and dataset files, set +`TRANSFORMERS_OFFLINE=0 HF_DATASETS_OFFLINE=0`. + +The output directory contains the run logs, generated DeepSpeed configs, +`loss_comparison.csv` with per-step losses, and `loss_comparison.txt` with the +maximum and mean absolute loss differences from the uninterrupted baseline. A +successful run means both resume paths produced comparable losses; inspect the +logs and differences when diagnosing a mismatch. + ## HumanEval results | Model | HumanEval (base) | HumanEval+ | diff --git a/training/deepspeed_finetune_demo/finetune_llama.py b/training/deepspeed_finetune_demo/finetune_llama.py index e3e826320..0f291c397 100644 --- a/training/deepspeed_finetune_demo/finetune_llama.py +++ b/training/deepspeed_finetune_demo/finetune_llama.py @@ -4,7 +4,7 @@ import argparse from datasets import load_dataset from torch.utils.data import DataLoader, DistributedSampler -from transformers import AutoModelForCausalLM, AutoTokenizer, default_data_collator +from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, default_data_collator from transformers.integrations.deepspeed import HfDeepSpeedConfig import json import random @@ -28,6 +28,15 @@ def set_seed(seed): DATASET_REGISTRY = { + "tatsu-lab/alpaca": { + "split": "train", + "preprocessor": "alpaca", + "field_map": { + "instruction": "instruction", + "input": "input", + "output": "output", + }, + }, "sahil2801/CodeAlpaca-20k": { "split": "train", "preprocessor": "alpaca", @@ -251,9 +260,81 @@ def _save_weights(model_engine, tokenizer, output_dir, step, keep_last=2): print_r(0, f"Saved checkpoint to {ckpt_dir}") +def _save_deepspeed_checkpoint(model_engine, output_dir, step, epoch, step_in_epoch): + """Save model, optimizer, scheduler, and resume position for exact continuation.""" + tag = f"step_{step}" + client_state = { + "global_step": step, + "epoch": epoch, + "step_in_epoch": step_in_epoch, + } + model_engine.save_checkpoint(output_dir, tag=tag, client_state=client_state) + print_r(0, f"Saved DeepSpeed checkpoint to {os.path.join(output_dir, tag)}") + + +def _load_deepspeed_checkpoint(model_engine, checkpoint_dir, tag): + load_path, client_state = model_engine.load_checkpoint(checkpoint_dir, tag=tag) + if load_path is None: + raise RuntimeError(f"DeepSpeed checkpoint was not loaded: {checkpoint_dir}/{tag}") + return client_state or {} + + +def _first_tensor(value): + if torch.is_tensor(value): + return value + if isinstance(value, (tuple, list)): + for item in value: + tensor = _first_tensor(item) + if tensor is not None: + return tensor + if isinstance(value, dict): + for item in value.values(): + tensor = _first_tensor(item) + if tensor is not None: + return tensor + return None + + +def _register_nonfinite_hooks(model): + state = {"found": False} + handles = [] + + def make_hook(name): + def hook(_module, _inputs, output): + if state["found"] or dist.get_rank() != 0: + return + tensor = _first_tensor(output) + if tensor is not None and not torch.isfinite(tensor).all().item(): + state["found"] = True + max_value = tensor.float().abs().max().item() + print_r(0, f"First non-finite module output: {name}, shape={tuple(tensor.shape)}, abs_max={max_value}") + return hook + + for name, module in model.named_modules(): + handles.append(module.register_forward_hook(make_hook(name))) + return handles + + +def _reset_rotary_embeddings(model): + for module in model.modules(): + rotary_emb = getattr(module, "rotary_emb", None) + if rotary_emb is None or not hasattr(rotary_emb, "inv_freq"): + continue + inv_freq = 1.0 / ( + rotary_emb.base + ** (torch.arange(0, rotary_emb.dim, 2, dtype=torch.float32) / rotary_emb.dim) + ) + rotary_emb.register_buffer("inv_freq", inv_freq, persistent=False) + rotary_emb.max_seq_len_cached = None + + def main(args): logging.basicConfig(level=logging.INFO, filename="pytorch_log.txt") set_seed(args.seed) + # Moonlight's checked-in modeling file imports this legacy helper. + from transformers.utils import import_utils + if not hasattr(import_utils, "is_torch_fx_available"): + import_utils.is_torch_fx_available = lambda: True # override batch size in ds_config with open(args.deepspeed_config, "r") as f: @@ -266,6 +347,10 @@ def main(args): tokenizer = AutoTokenizer.from_pretrained(args.model_name, trust_remote_code=True) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token + model_config = AutoConfig.from_pretrained(args.model_name, trust_remote_code=True) + if (isinstance(model_config.rope_scaling, dict) + and model_config.rope_scaling.get("rope_type") == "default"): + model_config.rope_scaling = None try: import flash_attn @@ -274,10 +359,12 @@ def main(args): attn_impl = None model = AutoModelForCausalLM.from_pretrained( args.model_name, + config=model_config, torch_dtype=torch.bfloat16, trust_remote_code=True, attn_implementation=attn_impl, ) + _reset_rotary_embeddings(model) model.config.use_cache = False model.gradient_checkpointing_enable() @@ -347,6 +434,8 @@ def main(args): train_sampler = DistributedSampler( tokenized_train_dataset, + num_replicas=dist.get_world_size(), + rank=dist.get_rank(), shuffle=True, seed=args.seed, ) @@ -360,7 +449,16 @@ def main(args): ) model_engine.train() + nonfinite_hooks = _register_nonfinite_hooks(model_engine.module) if args.debug_nonfinite else [] global_step = 0 + start_epoch = 0 + start_step_in_epoch = 0 + if args.resume_tag is not None: + client_state = _load_deepspeed_checkpoint(model_engine, args.resume_dir or args.output_dir, args.resume_tag) + global_step = int(client_state.get("global_step", 0)) + start_epoch = int(client_state.get("epoch", 0)) + start_step_in_epoch = int(client_state.get("step_in_epoch", 0)) + print_r(0, f"Resumed checkpoint at global step {global_step}") total_time = 0 total_count = 0 @@ -388,11 +486,13 @@ def main(args): wandb.init(project="deepspeed_finetune_demo", name=args.wandb_name) global_samples = 0 - for epoch in range(args.num_train_epochs): + for epoch in range(start_epoch, args.num_train_epochs): print_r(0, f"Starting epoch {epoch + 1}/{args.num_train_epochs}") train_dataloader.sampler.set_epoch(epoch) for step, batch in enumerate(train_dataloader): + if epoch == start_epoch and step < start_step_in_epoch: + continue if prof != None and global_step == args.profile_start: prof.start() if prof != None and global_step - args.profile_start == args.profile_steps: @@ -409,9 +509,18 @@ def main(args): batch = {k: v.to(model_engine.device) for k, v in batch.items()} outputs = model_engine(**batch) loss = outputs.loss + if global_step == 0 and dist.get_rank() == 0: + valid_labels = int((batch["labels"][:, :-1] != -100).sum().item()) + finite_params = all(torch.isfinite(param).all().item() for param in model_engine.module.parameters()) + finite_logits = torch.isfinite(outputs.logits).all().item() + logits_max = outputs.logits.float().abs().max().item() + print_r(0, f"First batch valid shifted labels: {valid_labels}, finite params: {finite_params}, " + f"finite logits: {finite_logits}, logits abs max: {logits_max:.4g}, " + f"finite loss: {torch.isfinite(loss).item()}") model_engine.backward(loss) model_engine.step() + global_step += 1 global_samples += model_engine.train_batch_size() step_time = time.time() - step_start_time @@ -460,8 +569,9 @@ def main(args): and global_step % args.checkpoint_steps == 0 and save_checkpoint_p ): - _save_weights(model_engine, tokenizer, args.output_dir, global_step) - global_step += 1 + if not args.skip_weight_export: + _save_weights(model_engine, tokenizer, args.output_dir, global_step) + _save_deepspeed_checkpoint(model_engine, args.output_dir, global_step, epoch, step + 1) if prof != None: prof.step() if args.max_steps > 0 and global_step >= args.max_steps: @@ -473,11 +583,16 @@ def main(args): if args.max_steps > 0 and global_step >= args.max_steps: break + for handle in nonfinite_hooks: + handle.remove() + if args.bench_start >= 0 and args.bench_steps > 0: print_r(0, f"Average iteration time = {total_time / total_count}") - if save_checkpoint_p: - _save_weights(model_engine, tokenizer, args.output_dir, global_step) + if save_checkpoint_p and args.save_final_checkpoint: + if not args.skip_weight_export: + _save_weights(model_engine, tokenizer, args.output_dir, global_step) + _save_deepspeed_checkpoint(model_engine, args.output_dir, global_step, epoch, step + 1) print_r(0, "Training complete!") @@ -520,6 +635,17 @@ def main(args): "--checkpoint_steps", type=int, default=0, help="Save a checkpoint every N steps (0 disables); keeps last 2", ) + parser.add_argument( + "--resume_tag", type=str, default=None, + help="DeepSpeed checkpoint tag to resume from", + ) + parser.add_argument( + "--resume_dir", type=str, default=None, + help="Directory containing resume_tag (defaults to output_dir)", + ) + parser.add_argument("--skip_weight_export", action="store_true") + parser.add_argument("--save_final_checkpoint", action="store_true") + parser.add_argument("--debug_nonfinite", action="store_true") parser.add_argument( "--eval_batch_size", type=int, default=4, help="Eval batch size per rank" ) diff --git a/training/deepspeed_finetune_demo/run_autoep_affine_ir_checkpoint_experiment.sh b/training/deepspeed_finetune_demo/run_autoep_affine_ir_checkpoint_experiment.sh new file mode 100755 index 000000000..1a8defc3f --- /dev/null +++ b/training/deepspeed_finetune_demo/run_autoep_affine_ir_checkpoint_experiment.sh @@ -0,0 +1,186 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team + +set -euo pipefail + +# Run a 50-step checkpoint/resume comparison for Moonlight AutoEP. +# +# The script compares: +# 1. uninterrupted training to step 100; +# 2. native DeepSpeed checkpoint resume at step 50; +# 3. native checkpoint -> universal conversion -> universal resume at step 50. +# +# Override MODEL_NAME, DATASET_NAME, NUM_GPUS, and OUTPUT_ROOT as needed. + +NUM_GPUS="${NUM_GPUS:-8}" +AUTOEP_SIZE="${AUTOEP_SIZE:-${NUM_GPUS}}" +ZERO_STAGE="${ZERO_STAGE:-2}" +GPU_INCLUDE="${GPU_INCLUDE:-}" +MODEL_NAME="${MODEL_NAME:-moonshotai/Moonlight-16B-A3B}" +DATASET_NAME="${DATASET_NAME:-tatsu-lab/alpaca}" +OUTPUT_ROOT="${OUTPUT_ROOT:-${PWD}/autoep_affine_ir_experiment}" +HF_HOME="${HF_HOME:-}" +DEEPSPEED_REPO="${DEEPSPEED_REPO:-}" +EXAMPLE_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +CONFIG="${EXAMPLE_DIR}/configs/z2_moonlight_autoep_adam.json" +TRAIN="${EXAMPLE_DIR}/finetune_llama.py" +CONVERTER="${DEEPSPEED_REPO}/deepspeed/checkpoint/ds_to_universal.py" +LAUNCHER="${DEEPSPEED_LAUNCHER:-ds}" +BASE_CONFIG="${OUTPUT_ROOT}/autoep_config.json" +UNIVERSAL_CONFIG="${OUTPUT_ROOT}/autoep_universal_config.json" +COMMON_ARGS=( + --model_name "${MODEL_NAME}" + --dataset_name "${DATASET_NAME}" + --batch_size 128 + --max_length 64 + --num_train_epochs 20 + --seed 42 + --skip_weight_export +) + +export HF_HOME +export AUTOEP_SIZE +export ZERO_STAGE +export TRANSFORMERS_OFFLINE="${TRANSFORMERS_OFFLINE:-1}" +export HF_DATASETS_OFFLINE="${HF_DATASETS_OFFLINE:-1}" +export PYTHONPATH="${DEEPSPEED_REPO}:${PYTHONPATH:-}" +export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}" +export TRITON_CACHE_DIR="${TRITON_CACHE_DIR:-/tmp/triton_cache_${USER:-copilot}}" +mkdir -p "${TRITON_CACHE_DIR}" + +mkdir -p "${OUTPUT_ROOT}" +cp "${CONFIG}" "${BASE_CONFIG}" + +python - "${BASE_CONFIG}" "${UNIVERSAL_CONFIG}" <<'PY' +import json +import os +import sys + +source, target = sys.argv[1:3] +with open(source) as handle: + config = json.load(handle) +config["train_batch_size"] = 128 +config["gradient_accumulation_steps"] = 16 +config["zero_optimization"]["stage"] = int(os.environ["ZERO_STAGE"]) +config["zero_optimization"]["offload_optimizer"] = {"device": "cpu", "pin_memory": False} +if config["zero_optimization"]["stage"] == 3: + config["zero_optimization"]["offload_param"] = {"device": "cpu", "pin_memory": False} +config["expert_parallel"]["autoep_size"] = int(os.environ["AUTOEP_SIZE"]) +with open(source, "w") as handle: + json.dump(config, handle, indent=2) + handle.write("\n") +config["checkpoint"] = {"load_universal": True} +with open(target, "w") as handle: + json.dump(config, handle, indent=2) + handle.write("\n") +PY + +run_train() { + local output_dir="$1" + local config="$2" + shift 2 + local launcher_args=(--num_gpus="${NUM_GPUS}") + if [[ -n "${GPU_INCLUDE}" ]]; then + launcher_args=(--include="${GPU_INCLUDE}") + fi + "${LAUNCHER}" "${launcher_args[@]}" "${TRAIN}" \ + "${COMMON_ARGS[@]}" \ + --output_dir "${output_dir}" \ + --deepspeed_config "${config}" \ + --checkpoint_steps 50 \ + "$@" +} + +echo "Running uninterrupted 100-step baseline..." +run_train "${OUTPUT_ROOT}/baseline_100" "${BASE_CONFIG}" --max_steps 100 \ + --checkpoint_steps 0 \ + > "${OUTPUT_ROOT}/baseline_100.log" 2>&1 + +echo "Running native checkpoint to step 50..." +run_train "${OUTPUT_ROOT}/native_50" "${BASE_CONFIG}" --max_steps 50 \ + > "${OUTPUT_ROOT}/native_50.log" 2>&1 + +echo "Resuming from native checkpoint to step 100..." +run_train "${OUTPUT_ROOT}/native_resume_100" "${BASE_CONFIG}" --max_steps 100 \ + --resume_dir "${OUTPUT_ROOT}/native_50" --resume_tag step_50 \ + --checkpoint_steps 0 \ + > "${OUTPUT_ROOT}/native_resume_100.log" 2>&1 + +echo "Converting step 50 checkpoint to universal format..." +python "${CONVERTER}" \ + --input_folder "${OUTPUT_ROOT}/native_50/step_50" \ + --output_folder "${OUTPUT_ROOT}/native_50/step_50_universal" \ + --num_extract_workers 1 \ + --num_merge_workers 1 \ + --no_strict \ + > "${OUTPUT_ROOT}/convert_universal.log" 2>&1 +expert_state_files=("${OUTPUT_ROOT}/native_50/step_50"/layer_*_expert_*_model_states.pt) +if [[ -e "${expert_state_files[0]}" ]]; then + cp "${expert_state_files[@]}" "${OUTPUT_ROOT}/native_50/step_50_universal/" +fi + +echo "Resuming from universal checkpoint to step 100..." +run_train "${OUTPUT_ROOT}/universal_resume_100" "${UNIVERSAL_CONFIG}" --max_steps 100 \ + --resume_dir "${OUTPUT_ROOT}/native_50" --resume_tag step_50_universal \ + --checkpoint_steps 0 \ + > "${OUTPUT_ROOT}/universal_resume_100.log" 2>&1 + +python - "${OUTPUT_ROOT}" <<'PY' +import csv +import math +import os +import re +import sys + +root = sys.argv[1] +pattern = re.compile(r"Step (\d+), Loss: ([0-9eE.+-]+)") +logs = { + "baseline": os.path.join(root, "baseline_100.log"), + "native_resume": os.path.join(root, "native_resume_100.log"), + "universal_resume": os.path.join(root, "universal_resume_100.log"), +} + +losses = {} +for name, path in logs.items(): + values = {} + with open(path) as handle: + for line in handle: + match = pattern.search(line) + if match: + values[int(match.group(1))] = float(match.group(2)) + losses[name] = values + +steps = sorted(set.intersection(*(set(values) for values in losses.values()))) +steps = [step for step in steps if 51 <= step <= 100] +if not steps: + raise RuntimeError("No common loss entries were found for steps 51-100.") + +csv_path = os.path.join(root, "loss_comparison.csv") +with open(csv_path, "w", newline="") as handle: + writer = csv.writer(handle) + writer.writerow(["step", *losses]) + for step in steps: + writer.writerow([step, *(losses[name][step] for name in losses)]) + +summary_path = os.path.join(root, "loss_comparison.txt") +with open(summary_path, "w") as handle: + for name in ("native_resume", "universal_resume"): + errors = [ + abs(losses["baseline"][step] - losses[name][step]) + for step in steps + ] + max_error = max(errors) + mean_error = sum(errors) / len(errors) + handle.write( + f"{name}: steps={len(steps)} max_abs_error={max_error:.9g} " + f"mean_abs_error={mean_error:.9g}\n" + ) + handle.write(f"Compared steps: {steps[0]}-{steps[-1]}\n") + +print(f"Loss comparison written to {csv_path}") +print(open(summary_path).read(), end="") +PY + +echo "Experiment completed." +echo "Logs and comparison: ${OUTPUT_ROOT}/*.log, ${OUTPUT_ROOT}/loss_comparison.{csv,txt}"