From f74f9906b8237b29a67ce253e594ad3e435c4240 Mon Sep 17 00:00:00 2001 From: 0z5a Date: Wed, 14 Oct 2026 19:44:18 +0200 Subject: [PATCH 1/2] Add vocabulary-parallel KL support to OPSD Signed-off-by: 0z5a --- training/opsd/README.md | 119 +++++++ .../benchmarks/bench_vocab_parallel_opsd.py | 332 ++++++++++++++++++ training/opsd/config.py | 14 + .../opsd/configs/ds_vocab_parallel_smoke.json | 21 ++ .../configs/opsd_vocab_parallel_smoke.json | 40 +++ .../opsd/data/vocab_parallel_prompts.jsonl | 16 + training/opsd/losses.py | 47 ++- training/opsd/main.py | 49 ++- training/opsd/teacher.py | 51 ++- .../distributed/run_numeric_and_preflight.py | 138 ++++++++ .../distributed/run_trainer_integration.py | 228 ++++++++++++ .../distributed/run_vocab_parallel_probe.py | 213 +++++++++++ training/opsd/tests/test_vocab_parallel.py | 106 ++++++ training/opsd/trainer.py | 59 ++-- training/opsd/vocab_parallel.py | 179 ++++++++++ 15 files changed, 1563 insertions(+), 49 deletions(-) create mode 100644 training/opsd/benchmarks/bench_vocab_parallel_opsd.py create mode 100644 training/opsd/configs/ds_vocab_parallel_smoke.json create mode 100644 training/opsd/configs/opsd_vocab_parallel_smoke.json create mode 100644 training/opsd/data/vocab_parallel_prompts.jsonl create mode 100644 training/opsd/tests/distributed/run_numeric_and_preflight.py create mode 100644 training/opsd/tests/distributed/run_trainer_integration.py create mode 100644 training/opsd/tests/distributed/run_vocab_parallel_probe.py create mode 100644 training/opsd/tests/test_vocab_parallel.py create mode 100644 training/opsd/vocab_parallel.py diff --git a/training/opsd/README.md b/training/opsd/README.md index e79e32b01..79a403229 100644 --- a/training/opsd/README.md +++ b/training/opsd/README.md @@ -118,6 +118,122 @@ python -m pytest tests/ -v See `configs/opsd_hybrid_engine.json` for a fully-populated example. +### Vocabulary-parallel KL + +The default `distillation.loss_backend` is `"dense"`. To keep training and +teacher-cache logits in their actual AutoTP vocabulary shards, opt in with: + +```json +{ + "teacher": {"autotp_size": 2, "offload_to_cpu": false}, + "distillation": { + "loss_backend": "vocab_parallel", + "loss_type": "reverse_kl", + "temperature": 1.0, + "chunk_size": 32 + }, + "rollout": { + "temperature": 0.0, + "n_samples_per_prompt": 1, + "use_graph_capture": false, + "use_shared_prefill": false + } +} +``` + +These fields supplement the existing model, data and training settings. +The DeepSpeed JSON must set the same `tensor_parallel.autotp_size` for the +student and select student ZeRO-0 or ZeRO-2. The entry point enables +`tensor_parallel.vocab_parallel_lm_head` before Engine initialization. +Vocabulary-parallel KL requires the DeepSpeed operators introduced by +[PR #8791](https://github.com/deepspeedai/DeepSpeed/pull/8791). +The dense default explicitly retains full-vocabulary heads on both models, +including tied heads that a HuggingFace TP plan would otherwise shard. + +The adapter checks each producer's actual shard ranges, ordered TP members, +physical vocabulary width and tokenizer token-ID semantics. TP peers must +consume identical tokens and masks; DP replicas may use different samples +and lengths. TP=1 keeps the vocabulary local even when DP>1. Teacher logits +move directly from the local shard to a host cache; they are not gathered +and sliced. The cache completes D2H before becoming readable and must not +be used after `free()`. + +Greedy eager rollout temporarily gathers only the last position's vocabulary +logits, so generation selects a global token. Training and teacher forwards +retain local logits. The hook preserves the head's Parameter objects and +tied-weight identity, and is removed on both success and failure. The +existing next-token shift, temperature-squared scaling, local response-token +mean and Engine gradient-accumulation scaling remain unchanged. + +The supported opt-in route uses forward/reverse KL, single-sample greedy HF +generation, no teacher CPU offload, no pipeline/sequence parallelism and no +compiled or accelerated HybridEngine training. JSD, stochastic rollout, +graph capture and shared prefill are rejected. Frozen Qwen3 teachers can +fail in the dependency's replicated-parameter AutoTP gradient-hook setup; +that initialization issue is separate from the KL adapter. + +Run the device integration harnesses from this directory with its modules +on `PYTHONPATH`: + +```sh +PYTHONPATH="$PWD" torchrun --standalone --nproc_per_node=2 \ + tests/distributed/run_numeric_and_preflight.py --out artifacts/numeric +PYTHONPATH="$PWD" torchrun --standalone --nproc_per_node=4 \ + tests/distributed/run_trainer_integration.py --tp 2 --dtype bfloat16 \ + --tied --student-zero 2 --vocab-size 67 --empty-mask \ + --phase continuous --checkpoint artifacts/checkpoint --out artifacts/trainer +``` + +The numeric harness compares the actual CUDA operator with an FP64 dense +oracle and checks uniform rejection of malformed layouts and inconsistent +TP inputs. The trainer harness uses random tiny Qwen2 models and real +Engine optimizers; these tests qualify integration mechanics, rather than +pretrained distillation quality. Its `save` phase performs two optimizer +updates and its `resume` phase loads the checkpoint in a new process before +performing the remaining three. Compare those final parameters with a +separate five-update `continuous` run. + +`benchmarks/bench_vocab_parallel_opsd.py` measures a full online training +update or replay of frozen real rollout responses. `upstream_dense` imports +an unchanged upstream OPSD directory supplied by `--upstream-opsd-dir`; +the harness explicitly selects the dependency's documented full-logits +configuration for both producers, including the upstream teacher initializer. +`gathered_reference` uses the same candidate head layout with dense math and +a local-slice gather backward; `vocab_parallel` retains local logits. Replay +excludes generation. Online always generates new responses. Report actual +optimizer updates, unique response tokens counted once per TP group, each +run's workload hash, cache payload and device peak; distinguish these from +the original trainer's unsynchronized diagnostic phase timers. + +The Qwen2.5 smoke config uses the same 16 math/instruction prompts as the +recorded short-input runs. Replace the model IDs with pinned local snapshot +paths when reproducing a particular result. The measured snapshots were +`7ae557604adf67be50417f59c2c2f167def9a775` for the student and +`989aa7980e4cf806f80c7fef2b1adb7bc71aa306` for the teacher. + +```sh +PYTHONPATH="$PWD" torchrun --standalone --nproc_per_node=2 \ + benchmarks/bench_vocab_parallel_opsd.py \ + --config configs/opsd_vocab_parallel_smoke.json \ + --mode online --implementation vocab_parallel --sdpa-backend math \ + --fixture artifacts/smoke-fixture --warmup-updates 0 --measure-updates 5 \ + --output-jsonl artifacts/smoke.jsonl +``` + +The benchmark checks the actual FP32 optimizer parameters and rejects any +skipped update. `Engine.global_steps` alone is insufficient, since overflow +can advance that counter without updating weights. Model-weight hashes are +reported separately from optimizer-master hashes. + +`--sdpa-backend math` selects PyTorch's public math SDPA backend uniformly +for A, B and C. With Torch 2.12.1+cu130 on H100, the default SDPA selection +produced finite logits and KL logit gradients but NaN model gradients on a +long-input, early-EOS batch; the optimizer skipped the update. The matched +math-backend diagnostic had finite gradients and changed FP32 masters. +This is a measured backend limitation, separate from vocabulary-parallel +KL. The benchmark does not change the application's default attention +backend, and these measurements do not qualify every PyTorch/backend pair. + ## Design notes * **Why CPU-cache the teacher logits?** Holding both student and teacher @@ -126,6 +242,9 @@ See `configs/opsd_hybrid_engine.json` for a fully-populated example. the worst-case GPU footprint of the loss path. The streamed loss (`losses.streamed_distillation_loss`) pulls teacher chunks back to GPU one sequence slice at a time so the full tensor never re-materialises. + A teacher forward still materializes its output, and backward may retain + tensors from multiple loss chunks. Chunk size bounds a copy/temporary + workspace; it does not establish a bound on full-step device peak memory. * **Why an abstract `RolloutEngine`?** The ABC keeps the trainer engine-agnostic so additional backends can be added without touching the diff --git a/training/opsd/benchmarks/bench_vocab_parallel_opsd.py b/training/opsd/benchmarks/bench_vocab_parallel_opsd.py new file mode 100644 index 000000000..aac4bb8eb --- /dev/null +++ b/training/opsd/benchmarks/bench_vocab_parallel_opsd.py @@ -0,0 +1,332 @@ +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +"""Measure actual online OPSD or replay of previously generated responses. + +The upstream baseline imports a separately frozen, unchanged Examples tree. +The gathered reference retains candidate head layouts, gathers teacher +logits before D2H, and uses dense KL with local-slice gather backward. +""" + +import argparse +import hashlib +import json +import sys +import time +from pathlib import Path + +import torch +import torch.nn.functional as F +from deepspeed import comm as dist +from deepspeed.utils import groups, safe_get_full_fp32_param + + +def full_logits(local, layout): + if len(layout.ranks) == 1: + return local + widths = [end - start for start, end in layout.ranges] + padded = F.pad(local.contiguous(), (0, max(widths) - local.shape[-1])) + values = [torch.empty_like(padded) for _ in widths] + dist.all_gather(values, padded, group=layout.group) + return torch.cat([value[..., :width] for value, width in zip(values, widths)], -1) + + +class DenseGather(torch.autograd.Function): + @staticmethod + def forward(ctx, local, layout): + ctx.start, ctx.end = layout.start, layout.end + return full_logits(local, layout) + + @staticmethod + def backward(ctx, gradient): + return gradient[..., ctx.start:ctx.end].contiguous(), None + + +def run(): + parser = argparse.ArgumentParser() + parser.add_argument('--mode', choices=['online', 'replay'], required=True) + parser.add_argument('--implementation', choices=['upstream_dense', 'dense_default', 'gathered_reference', 'vocab_parallel'], + required=True) + parser.add_argument('--config', required=True) + parser.add_argument('--fixture', required=True, help='Per-rank real-rollout fixture directory') + parser.add_argument('--upstream-opsd-dir') + parser.add_argument('--warmup-updates', type=int, default=5) + parser.add_argument('--measure-updates', type=int, default=30) + parser.add_argument('--output-jsonl', required=True) + parser.add_argument('--save-checkpoint') + parser.add_argument('--sdpa-backend', choices=['default', 'math'], default='default', + help='Use the same PyTorch SDPA backend for every measured implementation') + parser.add_argument('--profile', action='store_true', help='Additional phase synchronization; separate from throughput') + args = parser.parse_args() + if args.implementation == 'upstream_dense': + if not args.upstream_opsd_dir: + parser.error('upstream_dense needs the unchanged frozen upstream OPSD directory') + sys.path.insert(0, args.upstream_opsd_dir) + else: + sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + import main as application + from teacher import TeacherLogitCache + from losses import streamed_distillation_loss + initialization_overrides = [] + original_initialize = application.deepspeed.initialize + if args.implementation == 'upstream_dense': + def full_vocabulary_initialize(*parameters, **keywords): + config = keywords.get('config') + if isinstance(config, dict) and config.get('tensor_parallel', {}).get('autotp_size', 1) > 1: + config = dict(config) + config['tensor_parallel'] = dict(config['tensor_parallel'], vocab_parallel_lm_head=False) + keywords['config'] = config + initialization_overrides.append(config['tensor_parallel']) + return original_initialize(*parameters, **keywords) + # The unchanged upstream TeacherWrapper does not expose a head flag. + # Set the documented full-logits configuration on both initializers; + # preserve its model, cache, loss, backward and optimizer code. + application.deepspeed.initialize = full_vocabulary_initialize + + class MeasuredTrainer(application.OPSDTrainer): + def __init__(self, *parameters, **keywords): + super().__init__(*parameters, **keywords) + self.rank = dist.get_rank() + self.tp_rank = groups.get_tensor_model_parallel_rank() + self.samples, self.fixture_rows, self.measured = [], [], [] + self.elapsed = 0.0 + self.tokens = 0 + self.initial_updates = self.student_engine.global_steps + self.fixture_path = Path(args.fixture) / ('rank-' + str(self.rank) + '.pt') + self.replay = torch.load(self.fixture_path, map_location='cpu', weights_only=True) if args.mode == 'replay' else None + self.initial = self._parameter_digest() + self.initial_optimizer = self._optimizer_digest() + self.initial_skipped = self.student_engine.skipped_steps + original_forward = self.teacher.forward_to_cache + self.raw_teacher_forward = original_forward + original_generate = self.rollout.generate + + def observed_generate(*parameters, **keywords): + result = original_generate(*parameters, **keywords) + self.last_rollout = result + return result + + self.rollout.generate = observed_generate + + def observed_teacher(input_ids, attention_mask): + if args.mode == 'online': + from utils import build_response_mask + mask = build_response_mask(self.last_rollout.response_start_idx.to(input_ids.device), attention_mask) + self.fixture_rows.append({'input_ids': input_ids.detach().cpu(), + 'attention_mask': attention_mask.detach().cpu(), 'response_mask': mask.detach().cpu(), + 'source_microstep': self.step}) + if args.profile: + torch.cuda.synchronize() + started = time.perf_counter() + if args.implementation == 'gathered_reference': + with torch.no_grad(): + logits = self.teacher._callable(input_ids=input_ids, attention_mask=attention_mask).logits + cache = TeacherLogitCache.from_gpu_logits(full_logits(logits, self.teacher.vocab_layout)) + else: + cache = original_forward(input_ids, attention_mask) + if args.profile: + torch.cuda.synchronize() + self.last_teacher = {'host_cache_bytes': cache.cpu_logits.numel() * cache.cpu_logits.element_size(), + 'cache_shape': list(cache.shape), 'teacher_profile_s': time.perf_counter() - started} + return cache + + self.teacher.forward_to_cache = observed_teacher + + @torch.no_grad() + def _fixed_response_loss(self, fixture): + self.student_engine.eval() + ids, attention, mask = [fixture[name].to(self.device) + for name in ['input_ids', 'attention_mask', 'response_mask']] + logits = self.student_engine(input_ids=ids, attention_mask=attention).logits[:, :-1] + cache = self.raw_teacher_forward(ids, attention) + if args.implementation == 'gathered_reference': + logits = full_logits(logits, self.vocab_layout) + target = full_logits(cache.cpu_logits.to(logits.device), self.teacher.vocab_layout) + cache.free() + cache = TeacherLogitCache.from_gpu_logits(target) + backend = {} + if args.implementation == 'vocab_parallel': + backend = {'loss_backend': 'vocab_parallel', 'vocab_layout': self.vocab_layout} + try: + loss = streamed_distillation_loss(logits, + lambda start, end: cache.chunk_to_device(start, end, logits.device, logits.dtype), mask[:, 1:], + loss_type=self.cfg.distillation.loss_type, temperature=self.cfg.distillation.temperature, + chunk_size=self.cfg.distillation.chunk_size, **backend) + finally: + cache.free() + value = loss.detach().clone() + dist.all_reduce(value) + value /= dist.get_world_size() + assert torch.isfinite(value) + return float(value) + + def _parameter_digest(self): + digest = hashlib.sha256() + for name, parameter in self.student_engine.module.named_parameters(): + value = safe_get_full_fp32_param(parameter) + if value is None: + value = parameter.detach().float() + digest.update(name.encode()) + digest.update(value.detach().cpu().contiguous().numpy().tobytes()) + return digest.hexdigest() + + def _optimizer_digest(self): + # Unfused low-precision optimizers do not expose their masters + # through safe_get_full_fp32_param; public param_groups contain + # the actual FP32 parameters passed to the underlying optimizer. + digest = hashlib.sha256() + for group_index, group in enumerate(self.student_engine.optimizer.param_groups): + for parameter_index, parameter in enumerate(group['params']): + if parameter.dtype != torch.float32: + raise RuntimeError('This benchmark requires FP32 optimizer master parameters') + digest.update(str((group_index, parameter_index, tuple(parameter.shape))).encode()) + digest.update(parameter.detach().cpu().contiguous().numpy().tobytes()) + return digest.hexdigest() + + def _student_update(self, input_ids, attention_mask, response_mask, cache): + if args.implementation in ('vocab_parallel', 'dense_default'): + return super()._student_update(input_ids, attention_mask, response_mask, cache) + started = time.perf_counter() + self.student_engine.train() + local = self.student_engine(input_ids=input_ids, attention_mask=attention_mask).logits + if args.implementation == 'gathered_reference': + layout = self.vocab_layout + assert local.shape[-1] == layout.local_size and cache.shape == (*local.shape[:2], layout.vocab_size) + local = DenseGather.apply(local, layout) + # Replay of A uses the unchanged upstream loss function with its + # original full logits, next-token shift and Engine backward. + logits = local[:, :-1] + mask = response_mask[:, 1:] + loss = streamed_distillation_loss(logits, + lambda start, end: cache.chunk_to_device(start, end, logits.device, logits.dtype), mask, + loss_type=self.cfg.distillation.loss_type, temperature=self.cfg.distillation.temperature, + chunk_size=self.cfg.distillation.chunk_size) + self.student_engine.backward(loss) + self.student_engine.step() + return loss, int(mask.sum().item()), time.perf_counter() - started + + def _train_step(self, batch): + gas = self.cfg.training.gradient_accumulation_steps + if self.step % gas == 0: + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + self.elapsed, self.tokens = 0.0, 0 + torch.cuda.synchronize() + started = time.perf_counter() + if args.mode == 'online': + metrics = super()._train_step(batch) + else: + fixture = self.replay[self.step] + ids, attention, mask = [fixture[name].to(self.device) + for name in ['input_ids', 'attention_mask', 'response_mask']] + if getattr(self, 'vocab_layout', None) is not None: + from vocab_parallel import validate_tp_batch + validate_tp_batch(self.vocab_layout, ids, attention, mask) + cache = self.teacher.forward_to_cache(ids, attention) + try: + loss, tokens, student_s = self._student_update(ids, attention, mask, cache) + finally: + cache.free() + averaged = loss.detach().clone() + dist.all_reduce(averaged) + averaged /= dist.get_world_size() + metrics = {'loss': float(averaged.item()), 'response_tokens': tokens, 'rollout_s': 0.0, + 'teacher_s': self.last_teacher['teacher_profile_s'], 'student_s': student_s, 'step_s': 0.0} + torch.cuda.synchronize() + self.elapsed += time.perf_counter() - started + self.tokens += metrics['response_tokens'] if self.tp_rank == 0 else 0 + assert torch.isfinite(torch.tensor(metrics['loss'])) + if (self.step + 1) % gas == 0: + assert self.student_engine.skipped_steps == self.initial_skipped, 'Optimizer skipped an update' + update = self.student_engine.global_steps - self.initial_updates + row = {'rank': self.rank, 'mode': args.mode, 'implementation': args.implementation, + 'update': update, 'optimizer_skipped_steps': self.student_engine.skipped_steps - self.initial_skipped, + 'sdpa_backend': args.sdpa_backend, 'loss': metrics['loss'], 'elapsed_s': self.elapsed, + 'unique_response_tokens': self.tokens, + 'peak_allocated_bytes': torch.cuda.max_memory_allocated(), + 'peak_reserved_bytes': torch.cuda.max_memory_reserved(), **self.last_teacher} + ranks = [None] * dist.get_world_size() + dist.all_gather_object(ranks, row) + if update > args.warmup_updates: + self.measured.append(row) + if self.rank == 0: + elapsed = max(item['elapsed_s'] for item in ranks) + tokens = sum(item['unique_response_tokens'] for item in ranks) + summary = {'update': update, 'elapsed_s': elapsed, 'unique_response_tokens': tokens, + 'tokens_per_s': tokens / elapsed, 'ranks': ranks, + 'mode': args.mode, 'implementation': args.implementation, + 'profile': args.profile} + with Path(args.output_jsonl).open('a') as output: + output.write(json.dumps(summary) + '\n') + return metrics + + def train(self): + expected = args.warmup_updates + args.measure_updates + assert self.cfg.training.max_steps == expected * self.cfg.training.gradient_accumulation_steps + if args.mode == 'replay': + fixed_response = self.replay[0] + else: + from contextlib import nullcontext + from deepspeed.runtime.rollout import RolloutRequest, SamplingConfig + from utils import build_response_mask + batch = next(iter(self.dataloader)) + self.student_engine.eval() + layout = getattr(self, 'vocab_layout', None) + scope = nullcontext() + if layout is not None: + from vocab_parallel import gathered_rollout_logits + scope = gathered_rollout_logits(self.student_engine.module, layout) + with torch.no_grad(), scope: + response = self.rollout.generate(RolloutRequest(batch['prompt_ids'].to(self.device), + batch['prompt_attention_mask'].to(self.device)), + SamplingConfig(max_new_tokens=self.cfg.rollout.max_response_length, temperature=0.0)) + fixed_response = {'input_ids': response.input_ids.detach().cpu(), + 'attention_mask': response.attention_mask.detach().cpu(), + 'response_mask': build_response_mask(response.response_start_idx, + response.attention_mask).detach().cpu()} + fixed_response_initial = self._fixed_response_loss(fixed_response) + super().train() + assert self.student_engine.global_steps - self.initial_updates == expected + assert len(self.measured) == args.measure_updates + final = self._parameter_digest() + final_optimizer = self._optimizer_digest() + fixed_response_final = self._fixed_response_loss(fixed_response) + assert final_optimizer != self.initial_optimizer, 'No actual optimizer master update' + assert self.student_engine.skipped_steps == self.initial_skipped + assert all(not p.requires_grad and p.grad is None for p in self.teacher._callable.module.parameters()) + if args.mode == 'online': + self.fixture_path.parent.mkdir(parents=True, exist_ok=True) + torch.save(self.fixture_rows, self.fixture_path) + if args.save_checkpoint: + self.student_engine.save_checkpoint(args.save_checkpoint, tag='completed', + client_state={'opsd_microstep': self.step, 'optimizer_updates': expected}) + resolved = {'status': 'PASS', 'rank': self.rank, 'mode': args.mode, 'implementation': args.implementation, + 'optimizer_updates': expected, 'measured_updates': len(self.measured), + 'initial_parameter_sha256': self.initial, 'final_parameter_sha256': final, + 'config': self.cfg.to_dict(), 'fixture_sha256': hashlib.sha256(self.fixture_path.read_bytes()).hexdigest(), + 'initial_optimizer_sha256': self.initial_optimizer, 'final_optimizer_sha256': final_optimizer, + 'optimizer_skipped_steps': self.student_engine.skipped_steps - self.initial_skipped, + 'model_weights_changed': final != self.initial, 'sdpa_backend': args.sdpa_backend, + 'optimizer_class': type(self.student_engine.optimizer).__name__, + 'teacher_frozen': True, 'source_dir': str(Path(application.__file__).parent)} + resolved.update(fixed_response_loss_before=fixed_response_initial, fixed_response_loss_after=fixed_response_final, + fixed_response_scope='Initial response reused after training; separate from timed workload, no convergence claim', + full_vocabulary_initialization_overrides=initialization_overrides) + Path(args.output_jsonl + '.rank-' + str(self.rank) + '.json').write_text(json.dumps(resolved, indent=2)) + + application.OPSDTrainer = MeasuredTrainer + sys.argv = ['main.py', '--config', args.config] + try: + if args.sdpa_backend == 'math': + from torch.nn.attention import SDPBackend, sdpa_kernel + with sdpa_kernel(SDPBackend.MATH): + application.main() + else: + application.main() + finally: + application.deepspeed.initialize = original_initialize + if dist.is_initialized(): + dist.destroy_process_group() + + +if __name__ == '__main__': + run() diff --git a/training/opsd/config.py b/training/opsd/config.py index 2b75edb50..24a642c15 100644 --- a/training/opsd/config.py +++ b/training/opsd/config.py @@ -8,6 +8,7 @@ """ import json +import math from dataclasses import dataclass, field from typing import Optional @@ -17,6 +18,7 @@ @dataclass class RolloutConfig(_BaseRolloutConfig): """Extends DeepSpeed's RolloutConfig with OPSD generation knobs.""" + use_shared_prefill: bool = False max_prompt_length: int = 1024 max_response_length: int = 1024 temperature: float = 0.0 @@ -47,6 +49,7 @@ class DistillationConfig: loss_type: str = "reverse_kl" temperature: float = 1.0 chunk_size: int = 512 + loss_backend: str = "dense" @dataclass @@ -111,3 +114,14 @@ def validate(self) -> None: raise ValueError(f"Unknown rollout engine {self.rollout.engine!r}; expected 'hybrid_engine'") if self.distillation.chunk_size <= 0: raise ValueError("distillation.chunk_size must be positive") + if not math.isfinite(self.distillation.temperature) or self.distillation.temperature <= 0: + raise ValueError("distillation.temperature must be finite and positive") + if self.distillation.loss_backend not in ("dense", "vocab_parallel"): + raise ValueError(f"Unknown loss_backend {self.distillation.loss_backend!r}") + if self.distillation.loss_backend == "vocab_parallel": + if self.distillation.loss_type == "jsd": + raise ValueError("vocab_parallel supports forward/reverse KL; use dense for JSD") + if self.rollout.temperature != 0 or self.rollout.n_samples_per_prompt != 1: + raise ValueError("vocab_parallel currently supports single-sample greedy rollout") + if self.rollout.use_graph_capture or self.rollout.use_shared_prefill: + raise ValueError("vocab_parallel requires eager rollout without shared prefill") diff --git a/training/opsd/configs/ds_vocab_parallel_smoke.json b/training/opsd/configs/ds_vocab_parallel_smoke.json new file mode 100644 index 000000000..a78f1338c --- /dev/null +++ b/training/opsd/configs/ds_vocab_parallel_smoke.json @@ -0,0 +1,21 @@ +{ + "train_micro_batch_size_per_gpu": 1, + "gradient_accumulation_steps": 2, + "tensor_parallel": { + "autotp_size": 2, + "vocab_parallel_lm_head": true + }, + "zero_optimization": { + "stage": 0 + }, + "bf16": { + "enabled": true + }, + "gradient_clipping": 0.0, + "optimizer": { + "type": "SGD", + "params": { + "lr": 5e-06 + } + } +} diff --git a/training/opsd/configs/opsd_vocab_parallel_smoke.json b/training/opsd/configs/opsd_vocab_parallel_smoke.json new file mode 100644 index 000000000..fda4b6584 --- /dev/null +++ b/training/opsd/configs/opsd_vocab_parallel_smoke.json @@ -0,0 +1,40 @@ +{ + "student": { + "model_name_or_path": "Qwen/Qwen2.5-0.5B-Instruct", + "dtype": "bfloat16" + }, + "teacher": { + "model_name_or_path": "Qwen/Qwen2.5-1.5B-Instruct", + "dtype": "bfloat16", + "offload_to_cpu": false, + "autotp_size": 2 + }, + "distillation": { + "loss_backend": "vocab_parallel", + "loss_type": "reverse_kl", + "temperature": 1.0, + "chunk_size": 32 + }, + "rollout": { + "max_prompt_length": 128, + "max_response_length": 64, + "temperature": 0.0, + "use_graph_capture": false + }, + "training": { + "micro_batch_size_per_gpu": 1, + "gradient_accumulation_steps": 2, + "max_steps": 10, + "num_train_epochs": 20, + "learning_rate": 5e-06, + "logging_steps": 2, + "save_steps": 100000, + "seed": 42, + "save_dir": "artifacts/checkpoints" + }, + "data": { + "path": "data/vocab_parallel_prompts.jsonl", + "shuffle": false + }, + "deepspeed_config": "configs/ds_vocab_parallel_smoke.json" +} diff --git a/training/opsd/data/vocab_parallel_prompts.jsonl b/training/opsd/data/vocab_parallel_prompts.jsonl new file mode 100644 index 000000000..a5c981095 --- /dev/null +++ b/training/opsd/data/vocab_parallel_prompts.jsonl @@ -0,0 +1,16 @@ +{"prompt": "A box holds 12 red balls and 8 blue balls. How many balls are in the box? Explain the addition."} +{"prompt": "A train travels 60 kilometers in one hour. How far does it travel in three hours at the same speed?"} +{"prompt": "Anna has 25 euros. She buys a book costing 9 euros and a pen costing 3 euros. How much remains?"} +{"prompt": "Explain why the sum of two even integers is even, using an algebraic argument."} +{"prompt": "Compute the area of a rectangle that is 7 centimeters long and 4 centimeters wide."} +{"prompt": "If x plus 5 equals 17, find x and check your answer by substitution."} +{"prompt": "A recipe needs 200 grams of flour for four servings. How much flour is needed for six servings?"} +{"prompt": "A fair six-sided die is rolled once. What is the probability of getting an even number?"} +{"prompt": "Describe the water cycle in three concise sentences."} +{"prompt": "Write a Python function that returns the largest element of a nonempty list of integers."} +{"prompt": "Explain the difference between a Python list and a tuple, with one example of each."} +{"prompt": "Why does dividing a positive number by one half double its value? Give a numerical example."} +{"prompt": "Solve 3x minus 4 equals 11. Show the intermediate steps."} +{"prompt": "A shop discounts a 50 euro item by 20 percent. What is its sale price?"} +{"prompt": "Find the median of the numbers 8, 2, 5, 9, and 4."} +{"prompt": "Explain how to check whether a positive integer is prime."} diff --git a/training/opsd/losses.py b/training/opsd/losses.py index b9dba5369..033f7aad9 100644 --- a/training/opsd/losses.py +++ b/training/opsd/losses.py @@ -7,6 +7,8 @@ both student *and* teacher logits at once would double that. We chunk along the sequence axis so the per-chunk softmax + difference only ever needs ``[B, chunk, V]`` of working memory, regardless of T. +Backward can retain tensors from multiple chunks, so this temporary-workspace +bound does not describe the complete forward/backward device peak. Math conventions: * ``forward_kl`` = D_KL(teacher || student) — mode-covering for student @@ -61,6 +63,14 @@ def _jsd(student_logits: torch.Tensor, teacher_logits: torch.Tensor, temperature } +def require_vocab_parallel_kl(): + try: + from deepspeed.sequence.vocab_parallel_ops import vocab_parallel_kl_div + except ImportError as error: + raise RuntimeError("vocab_parallel requires DeepSpeed vocabulary-parallel KL ops (PR #8791)") from error + return vocab_parallel_kl_div + + def chunked_distillation_loss( student_logits: torch.Tensor, teacher_logits: torch.Tensor, @@ -126,6 +136,9 @@ def streamed_distillation_loss( loss_type: str = "reverse_kl", temperature: float = 1.0, chunk_size: int = 512, + *, + loss_backend: str = "dense", + vocab_layout=None, ) -> torch.Tensor: """Same math as :func:`chunked_distillation_loss`, but teacher logits are pulled chunk-by-chunk via a fetcher so the full ``[B, T, V]`` teacher @@ -144,6 +157,15 @@ def streamed_distillation_loss( if loss_type not in _LOSS_FNS: raise ValueError(f"Unknown loss_type {loss_type!r}; choose from {sorted(_LOSS_FNS)}") fn = _LOSS_FNS[loss_type] + if loss_backend not in ("dense", "vocab_parallel"): + raise ValueError(f"Unknown loss_backend {loss_backend!r}") + parallel_kl = None + if loss_backend == "vocab_parallel": + if loss_type == "jsd" or vocab_layout is None: + raise ValueError("vocab_parallel requires a validated layout and forward/reverse KL") + if student_logits.shape[-1] != vocab_layout.local_size: + raise ValueError("Student logits differ from the validated vocabulary shard") + parallel_kl = require_vocab_parallel_kl() B, T, _ = student_logits.shape if response_mask.shape != (B, T): @@ -152,22 +174,31 @@ def streamed_distillation_loss( mask_f = response_mask.to(torch.float32) total_tokens = mask_f.sum().clamp_min(1.0) - total_loss = student_logits.new_zeros((), dtype=torch.float32) + total_loss = student_logits[:, :0].sum(dtype=torch.float32) for start in range(0, T, chunk_size): end = min(start + chunk_size, T) chunk_mask = mask_f[:, start:end] - if chunk_mask.sum().item() == 0: + if parallel_kl is None and chunk_mask.sum().item() == 0: continue teacher_chunk = teacher_chunk_fetcher(start, end) - if teacher_chunk.shape[1] != (end - start): + if teacher_chunk.shape != student_logits[:, start:end].shape: raise RuntimeError(f"fetcher returned chunk of length {teacher_chunk.shape[1]}, " f"expected {end - start}") - per_tok = fn( - student_logits[:, start:end].float(), - teacher_chunk.float(), - temperature, - ) + if parallel_kl is None: + per_tok = fn(student_logits[:, start:end].float(), teacher_chunk.float(), temperature) + else: + student_chunk = student_logits[:, start:end].float() + teacher_chunk = teacher_chunk.detach().float() + if temperature != 1.0: + student_chunk = student_chunk / temperature + teacher_chunk = teacher_chunk / temperature + per_tok = parallel_kl(student_chunk, teacher_chunk, + reverse=loss_type == "reverse_kl", tp_group=vocab_layout.group) + per_tok = per_tok * temperature**2 + # Undefined padding distributions return NaN with zero gradients + # in PR-1; masking must happen before the weighted sum. + per_tok = per_tok.masked_fill(chunk_mask == 0, 0) total_loss = total_loss + (per_tok * chunk_mask).sum() return total_loss / total_tokens diff --git a/training/opsd/main.py b/training/opsd/main.py index dde605f30..6a2949ea3 100644 --- a/training/opsd/main.py +++ b/training/opsd/main.py @@ -28,6 +28,7 @@ from deepspeed.runtime.rollout import build_rollout from teacher import TeacherWrapper from trainer import OPSDTrainer +from losses import require_vocab_parallel_kl def _seed_everything(seed: int) -> None: @@ -47,6 +48,10 @@ def _load_ds_config(path: str) -> dict: return json.load(f) +def _feature_enabled(value): + return value.get("enabled", False) if isinstance(value, dict) else bool(value) + + def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--config", required=True, help="Path to OPSDConfig JSON") @@ -77,6 +82,27 @@ def main() -> None: ) ds_config = _load_ds_config(cfg.deepspeed_config) + vocab_parallel = cfg.distillation.loss_backend == "vocab_parallel" + if vocab_parallel: + require_vocab_parallel_kl() + tp = ds_config.setdefault("tensor_parallel", {}).get("autotp_size", 1) + if cfg.teacher.autotp_size != tp: + raise ValueError("vocab_parallel requires equal student and teacher AutoTP sizes") + if cfg.teacher.offload_to_cpu: + raise ValueError("vocab_parallel teacher CPU offload is not validated; disable offload_to_cpu") + if ds_config.get("zero_optimization", {}).get("stage", 0) not in (0, 2): + raise ValueError("vocab_parallel currently supports student ZeRO-0 and ZeRO-2") + if _feature_enabled(ds_config.get("hybrid_engine", {})) or _feature_enabled(ds_config.get("compile", {})): + raise ValueError("vocab_parallel requires eager HF rollout and uncompiled training") + if ds_config.get("pipeline", {}).get("pipeline_parallel_size", 1) != 1: + raise ValueError("vocab_parallel pipeline parallelism is not validated") + if ds_config.get("sequence_parallel_size", 1) != 1: + raise ValueError("vocab_parallel sequence parallelism is not validated") + ds_config["tensor_parallel"]["vocab_parallel_lm_head"] = True + elif ds_config.get("tensor_parallel", {}).get("autotp_size", 1) > 1: + if ds_config["tensor_parallel"].get("vocab_parallel_lm_head") is True: + raise ValueError("Dense distillation requires full vocabulary logits; select vocab_parallel for sharded heads") + ds_config["tensor_parallel"]["vocab_parallel_lm_head"] = False ds_config["train_micro_batch_size_per_gpu"] = cfg.training.micro_batch_size_per_gpu ds_config["gradient_accumulation_steps"] = cfg.training.gradient_accumulation_steps @@ -87,7 +113,20 @@ def main() -> None: ) # --- frozen teacher ---------------------------------------------------- - teacher = TeacherWrapper(cfg.teacher, world_size=dist_world_size()) + teacher = TeacherWrapper(cfg.teacher, world_size=dist_world_size(), vocab_parallel=vocab_parallel) + if vocab_parallel: + from vocab_parallel import prepare_vocab_parallel + + teacher_tokenizer = AutoTokenizer.from_pretrained(cfg.teacher.model_name_or_path, + trust_remote_code=cfg.teacher.trust_remote_code) + if teacher_tokenizer.pad_token is None: + teacher_tokenizer.pad_token = teacher_tokenizer.eos_token + student_layout, teacher.vocab_layout = prepare_vocab_parallel( + student_engine.module, teacher._callable.module, tokenizer, teacher_tokenizer) + student_engine._opsd_vocab_layout = student_layout + elif any(getattr(model.get_output_embeddings(), "is_vocab_parallel_lm_head", False) + for model in [student_engine.module, teacher._callable.module]): + raise ValueError("Dense distillation cannot consume vocabulary shards; select vocab_parallel or disable sharded heads") # --- rollout engine ---------------------------------------------------- rollout = build_rollout( @@ -113,15 +152,13 @@ def main() -> None: tp_size = tp_size or 1 # autotp_size() returns 0 when AutoTP is unconfigured dp_size = dist_world_size() // tp_size if dp_size > 1: - # TP groups are contiguous, so DP rank = global_rank // tp_size. - # Pass num_replicas/rank explicitly so the sampler shards only along - # the real DP dimension (not the full world, which would split data - # across TP ranks and deadlock the all-reduce). + # Use the engine's actual DP group so TP peers consume the same + # prompts without assuming a global-rank ordering. sampler = DistributedSampler( dataset, shuffle=cfg.data.shuffle, num_replicas=dp_size, - rank=deepspeed.comm.get_rank() // tp_size, + rank=deepspeed.comm.get_rank(student_engine.data_parallel_group), ) else: sampler = None diff --git a/training/opsd/teacher.py b/training/opsd/teacher.py index fb5ca9d5c..9afb29ec0 100644 --- a/training/opsd/teacher.py +++ b/training/opsd/teacher.py @@ -9,8 +9,9 @@ teacher weight buffers can be released before the student backward pass. 2. **Student phase.** Forward + backward on the student. The distillation loss pulls teacher logits back to GPU **one sequence chunk at a time** via - :meth:`TeacherLogitCache.chunk_to_device`, so peak GPU memory for teacher - data is only ``[B, chunk, V]``. + :meth:`TeacherLogitCache.chunk_to_device`. Each copy is bounded by + ``[B, chunk, V]``; the teacher forward and tensors retained by backward + must still be included when measuring full-step device memory. This module deliberately lazy-imports ``deepspeed`` and ``transformers`` so that the pure data-handling pieces (``TeacherLogitCache`` and the streamed @@ -18,7 +19,7 @@ not have a working DeepSpeed launcher. """ -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Optional, Tuple import torch @@ -38,6 +39,8 @@ class TeacherLogitCache: """ cpu_logits: torch.Tensor # [B, T, V] + layout: Optional[object] = None + _freed: bool = field(default=False, init=False, repr=False) def __post_init__(self) -> None: if self.cpu_logits.dim() != 3: @@ -46,30 +49,38 @@ def __post_init__(self) -> None: if self.cpu_logits.device.type != "cpu": raise ValueError(f"cpu_logits must live on CPU; got device " f"{self.cpu_logits.device}") + if self.layout is not None and self.cpu_logits.shape[-1] != self.layout.local_size: + raise ValueError("Teacher cache width differs from its vocabulary shard") @classmethod - def from_gpu_logits(cls, logits: torch.Tensor, store_dtype: torch.dtype = torch.bfloat16) -> "TeacherLogitCache": + def from_gpu_logits(cls, logits: torch.Tensor, store_dtype: torch.dtype = torch.bfloat16, + *, layout=None) -> "TeacherLogitCache": """Detach + downcast + move to (pinned) host memory. - ``non_blocking=True`` lets the copy overlap with the next CUDA op when - the destination is pinned; we try to pin and fall back silently if the - host doesn't support it (e.g. CPU-only test environments). + Complete the copy before exposing CPU storage. This also gives dense + and sharded paths the same copy-readiness contract. """ downcast = logits.detach().to(dtype=store_dtype) try: host = torch.empty(downcast.shape, dtype=store_dtype, pin_memory=True) - host.copy_(downcast, non_blocking=True) except RuntimeError: - host = downcast.cpu() - return cls(cpu_logits=host) + host = torch.empty(downcast.shape, dtype=store_dtype) + host.copy_(downcast, non_blocking=False) + return cls(cpu_logits=host, layout=layout) + + def _require_alive(self) -> None: + if self._freed: + raise RuntimeError("Teacher cache has been freed") @property def shape(self) -> Tuple[int, int, int]: + self._require_alive() s = self.cpu_logits.shape return (int(s[0]), int(s[1]), int(s[2])) @property def dtype(self) -> torch.dtype: + self._require_alive() return self.cpu_logits.dtype def chunk_to_device(self, @@ -86,13 +97,14 @@ def chunk_to_device(self, if not (0 <= start < end <= T): raise ValueError(f"chunk bounds [{start}, {end}) invalid for T={T}") chunk = self.cpu_logits[:, start:end] - out = chunk.to(device=device, dtype=dtype if dtype is not None else chunk.dtype, non_blocking=True) + out = chunk.to(device=device, dtype=dtype if dtype is not None else chunk.dtype, non_blocking=False) return out def free(self) -> None: """Drop the underlying buffer so a step's teacher cache can be GC'd before the next teacher forward.""" self.cpu_logits = torch.empty(0) + self._freed = True _DTYPE_MAP = { @@ -129,11 +141,13 @@ class TeacherWrapper: can fill them, which surfaces as a "size mismatch" load error. """ - def __init__(self, cfg: TeacherConfig, world_size: int): + def __init__(self, cfg: TeacherConfig, world_size: int, *, vocab_parallel: bool = False): import deepspeed from transformers import AutoModelForCausalLM self.cfg = cfg + self.vocab_parallel = vocab_parallel + self.vocab_layout = None dtype = _resolve_dtype(cfg.dtype) model = AutoModelForCausalLM.from_pretrained( @@ -164,8 +178,13 @@ def __init__(self, cfg: TeacherConfig, world_size: int): "fp16": {"enabled": dtype is torch.float16}, "zero_optimization": zero_opt, } - if use_autotp: - ds_config["tensor_parallel"] = {"autotp_size": cfg.autotp_size} + if use_autotp or vocab_parallel: + ds_config["tensor_parallel"] = { + "autotp_size": cfg.autotp_size, + # HF plans can infer vocabulary sharding for tied heads. + # Dense OPSD requires full logits from both producers. + "vocab_parallel_lm_head": vocab_parallel, + } self._callable, *_ = deepspeed.initialize(model=model, config=ds_config) @torch.no_grad() @@ -175,4 +194,6 @@ def forward_to_cache(self, store_dtype: torch.dtype = torch.bfloat16) -> TeacherLogitCache: """Run teacher forward and stage logits onto the host.""" outputs = self._callable(input_ids=input_ids, attention_mask=attention_mask) - return TeacherLogitCache.from_gpu_logits(outputs.logits, store_dtype=store_dtype) + if self.vocab_parallel and self.vocab_layout is None: + raise RuntimeError("Validate the teacher vocabulary layout before its first forward") + return TeacherLogitCache.from_gpu_logits(outputs.logits, store_dtype=store_dtype, layout=self.vocab_layout) diff --git a/training/opsd/tests/distributed/run_numeric_and_preflight.py b/training/opsd/tests/distributed/run_numeric_and_preflight.py new file mode 100644 index 000000000..22212b6fb --- /dev/null +++ b/training/opsd/tests/distributed/run_numeric_and_preflight.py @@ -0,0 +1,138 @@ +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +"""CUDA KL matrix and uniform distributed rejection before the first loss.""" +import argparse +import json +import os +from datetime import timedelta +from pathlib import Path +from types import SimpleNamespace + +import deepspeed +import deepspeed.comm as dist +import torch + +from losses import streamed_distillation_loss +from teacher import TeacherLogitCache +from vocab_parallel import VocabShardLayout, prepare_vocab_parallel, validate_tp_batch +from run_vocab_parallel_probe import make_model + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--out', required=True) + args = parser.parse_args() + torch.cuda.set_device(int(os.environ['LOCAL_RANK'])) + deepspeed.init_distributed(dist_backend='nccl', auto_mpi_discovery=False, timeout=timedelta(seconds=120)) + rank, world = dist.get_rank(), dist.get_world_size() + assert world == 2 + group = dist.new_group([0, 1]) + records = [] + for dtype in [torch.float32, torch.bfloat16]: + for vocab in [64, 67]: + sizes = [(vocab + 1) // 2, vocab // 2] + start = sum(sizes[:rank]);end = start + sizes[rank] + layout = VocabShardLayout(vocab, start, end, (0, 1), ((0, sizes[0]), (sizes[0], vocab)), group) + torch.manual_seed(710) + full = torch.randn(2, 7, vocab, device='cuda', dtype=dtype) + teacher = torch.randn_like(full) + mask = torch.tensor([[1, 0, 1, 1, 0, 1, 1], [0, 1, 1, 0, 1, 0, 1]], device='cuda') + for direction in ['forward_kl', 'reverse_kl']: + for tau in [0.5, 1.0, 2.0]: + for chunk in [1, 3, 12]: + student = full[..., start:end].clone().requires_grad_() + reference = full.double().detach().requires_grad_() + target = teacher.to(torch.bfloat16).double() + lp, lq = torch.log_softmax(reference / tau, -1), torch.log_softmax(target / tau, -1) + tokens = (lq.exp() * (lq - lp)).sum(-1) if direction == 'forward_kl' else (lp.exp() * (lp - lq)).sum(-1) + expected = (tokens * mask).sum() * tau**2 / mask.sum() + cache = TeacherLogitCache.from_gpu_logits(teacher[..., start:end], layout=layout) + actual = streamed_distillation_loss(student, + lambda a, b: cache.chunk_to_device(a, b, student.device, student.dtype), mask, + loss_type=direction, temperature=tau, chunk_size=chunk, + loss_backend='vocab_parallel', vocab_layout=layout) + actual.backward();expected.backward() + torch.testing.assert_close(actual.double(), expected, atol=1e-6, rtol=1e-5) + tolerance = (1e-6, 1e-5) if dtype == torch.float32 else (2e-4, 1e-2) + torch.testing.assert_close(student.grad.double(), reference.grad[..., start:end], + atol=tolerance[0], rtol=tolerance[1]) + cache.free() + records.append({'dtype': str(dtype), 'vocab': vocab, 'direction': direction, + 'temperature': tau, 'chunk': chunk, 'status': 'PASS'}) + # Undefined padded distributions must contribute neither NaN loss nor + # gradient. The independent oracle excludes padding before log-softmax. + torch.manual_seed(711) + full = torch.randn(2, 7, 67, device='cuda') + teacher = torch.randn_like(full) + mask = torch.ones(2, 7, device='cuda') + mask[1, 0] = 0 + full[1, 0] = -torch.inf + teacher[1, 0] = -torch.inf + start, end = (0, 34) if rank == 0 else (34, 67) + layout = VocabShardLayout(67, start, end, (0, 1), ((0, 34), (34, 67)), group) + student = full[..., start:end].clone().requires_grad_() + reference = full.double().detach().requires_grad_() + valid = mask.bool() + lp = torch.log_softmax(reference[valid], -1) + lq = torch.log_softmax(teacher.to(torch.bfloat16).double()[valid], -1) + expected = (lp.exp() * (lp - lq)).sum(-1).mean() + cache = TeacherLogitCache.from_gpu_logits(teacher[..., start:end], layout=layout) + actual = streamed_distillation_loss(student, lambda a, b: cache.chunk_to_device(a, b, student.device), + mask, chunk_size=3, loss_backend='vocab_parallel', vocab_layout=layout) + actual.backward();expected.backward() + torch.testing.assert_close(actual.double(), expected, atol=1e-6, rtol=1e-5) + torch.testing.assert_close(student.grad.double(), reference.grad[..., start:end], atol=1e-6, rtol=1e-5) + assert torch.count_nonzero(student.grad[1, 0]) == 0 + cache.free() + records.append({'case': 'masked_undefined_distribution', 'status': 'PASS'}) + torch.manual_seed(1781) + model = make_model(torch.float32, False, 32, 1) + engine, *_ = deepspeed.initialize(model=model, config={ + 'train_micro_batch_size_per_gpu': 1, 'tensor_parallel': {'autotp_size': 2, 'vocab_parallel_lm_head': True}}) + head = engine.module.get_output_embeddings() + original_end = head.vocab_end_index + if rank == 1: + head.vocab_end_index -= 1 + rejected = None + try: + VocabShardLayout.from_model(engine.module) + except ValueError as failure: + rejected = str(failure) + assert rejected and 'Vocabulary' in rejected + head.vocab_end_index = original_end + # Malformed local metadata must fail on both peers without a timeout. + if rank == 1: + del head.vocab_end_index + try: + VocabShardLayout.from_model(engine.module) + except ValueError: + pass + else: + raise AssertionError('Malformed head was accepted') + head.vocab_end_index = original_end + tokenizer = SimpleNamespace(get_vocab=lambda: {'a': 0, 'b': 1}, get_added_vocab=lambda: {}, pad_token_id=0) + different = SimpleNamespace(get_vocab=lambda: {'a': 1, 'b': 0} if rank else {'a': 0, 'b': 1}, + get_added_vocab=lambda: {}, pad_token_id=0) + try: + prepare_vocab_parallel(engine.module, engine.module, tokenizer, different) + except ValueError as failure: + assert 'semantics' in str(failure) + else: + raise AssertionError('Tokenizer mismatch was accepted') + layout = VocabShardLayout.from_model(engine.module) + try: + ids = torch.tensor([[rank + 1, 3]], device='cuda') + validate_tp_batch(layout, ids, torch.ones_like(ids), torch.ones_like(ids)) + except ValueError: + pass + else: + raise AssertionError('Different TP prompts were accepted') + out = Path(args.out);out.mkdir(parents=True, exist_ok=True) + (out / ('rank-' + str(rank) + '.json')).write_text(json.dumps({ + 'status': 'PASS', 'numeric_cases': records, 'uniform_negative_cases': 4, + 'scope': 'Actual CUDA PR1 KL plus real AutoTP head preflight; controlled errors, no watchdog success'}, indent=2)) + dist.barrier();dist.destroy_process_group() + + +if __name__ == '__main__': + main() diff --git a/training/opsd/tests/distributed/run_trainer_integration.py b/training/opsd/tests/distributed/run_trainer_integration.py new file mode 100644 index 000000000..36f3e6f79 --- /dev/null +++ b/training/opsd/tests/distributed/run_trainer_integration.py @@ -0,0 +1,228 @@ +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +"""Actual OPSD trainer, Engine optimizer, DP/TP and restart integration. + +The random tiny Qwen2 fixtures test application mechanics, not pretrained +distillation quality. The independent reference uses dense PyTorch KL on +the same physical TP layout, with local-slice gather backward semantics. +""" + +import argparse +import copy +import json +import os +from datetime import timedelta +from pathlib import Path + +import deepspeed +import numpy as np +import torch +import torch.nn.functional as F +from deepspeed import comm as dist +from deepspeed.runtime.rollout import build_rollout +from deepspeed.utils import safe_get_full_fp32_param, safe_get_full_grad +from tokenizers import Tokenizer, models, pre_tokenizers +from torch.utils.data import DataLoader, Subset +from torch.utils.data.distributed import DistributedSampler +from transformers import PreTrainedTokenizerFast + +from config import OPSDConfig, TeacherConfig +from data import LeftPaddedPromptCollator, PromptDataset +from teacher import TeacherWrapper +from trainer import OPSDTrainer +from vocab_parallel import VocabShardLayout, prepare_vocab_parallel +from run_vocab_parallel_probe import compare, make_model + + +def gather(logits, layout): + if len(layout.ranks) == 1: + return logits + widths = [end - start for start, end in layout.ranges] + padded = F.pad(logits.contiguous(), (0, max(widths) - logits.shape[-1])) + outputs = [torch.empty_like(padded) for _ in widths] + dist.all_gather(outputs, padded, group=layout.group) + return torch.cat([value[..., :width] for value, width in zip(outputs, widths)], -1) + + +class DenseGather(torch.autograd.Function): + @staticmethod + def forward(ctx, logits, layout): + ctx.start, ctx.end = layout.start, layout.end + return gather(logits, layout) + + @staticmethod + def backward(ctx, gradient): + return gradient[..., ctx.start:ctx.end].contiguous(), None + + +def master(parameter): + value = safe_get_full_fp32_param(parameter) + return parameter.detach().float() if value is None else value.detach().float() + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--out', required=True) + parser.add_argument('--tp', type=int, choices=[1, 2], required=True) + parser.add_argument('--dtype', choices=['float32', 'bfloat16'], default='float32') + parser.add_argument('--tied', action='store_true') + parser.add_argument('--phase', choices=['continuous', 'save', 'resume'], default='continuous') + parser.add_argument('--checkpoint', required=True) + parser.add_argument('--student-zero', type=int, choices=[0, 2], default=0) + parser.add_argument('--vocab-size', type=int, default=64) + parser.add_argument('--empty-mask', action='store_true') + args = parser.parse_args() + torch.cuda.set_device(int(os.environ['LOCAL_RANK'])) + deepspeed.init_distributed(dist_backend='nccl', auto_mpi_discovery=False, + timeout=timedelta(seconds=120)) + rank, world = dist.get_rank(), dist.get_world_size() + assert world % args.tp == 0 + out = Path(args.out) + out.mkdir(parents=True, exist_ok=True) + dtype = getattr(torch, args.dtype) + torch.manual_seed(1781) + student = make_model(dtype, args.tied, 32, 1, args.vocab_size) + reference = copy.deepcopy(student) + teacher_model = make_model(dtype, args.tied, 64, 2, args.vocab_size) + fixture = out / 'fixture' + if rank == 0: + fixture.mkdir() + teacher_model.save_pretrained(fixture / 'teacher') + tokens = ['[PAD]', '[EOS]', '[BOS]', '[UNK]'] + ['t' + str(i) for i in range(4, args.vocab_size)] + backend = Tokenizer(models.WordLevel(dict(zip(tokens, range(args.vocab_size))), unk_token='[UNK]')) + backend.pre_tokenizer = pre_tokenizers.Whitespace() + tokenizer = PreTrainedTokenizerFast(tokenizer_object=backend, pad_token='[PAD]', eos_token='[EOS]', + bos_token='[BOS]', unk_token='[UNK]', padding_side='left') + tokenizer.chat_template = "{% for message in messages %}{{ message['content'] }}{% endfor %}" + tokenizer.save_pretrained(fixture / 'tokenizer') + prompts = [' '.join('t' + str(4 + (i + j) % 24) for j in range(2 + i % 5)) for i in range(40)] + (fixture / 'prompts.jsonl').write_text(''.join(json.dumps({'prompt': p}) + '\n' for p in prompts)) + dist.barrier() + tokenizer = PreTrainedTokenizerFast.from_pretrained(fixture / 'tokenizer') + config = {'train_micro_batch_size_per_gpu': 1, 'gradient_accumulation_steps': 2, + 'tensor_parallel': {'autotp_size': args.tp, 'vocab_parallel_lm_head': True}, + 'zero_optimization': {'stage': args.student_zero}, 'gradient_clipping': 0.0, + 'zero_allow_untested_optimizer': True, + 'bf16': {'enabled': dtype == torch.bfloat16}, + 'optimizer': {'type': 'SGD', 'params': {'lr': 0.05, 'momentum': 0.9}}} + candidate, *_ = deepspeed.initialize(model=student, model_parameters=student.parameters(), config=config) + dense, *_ = deepspeed.initialize(model=reference, model_parameters=reference.parameters(), config=config) + teacher = TeacherWrapper(TeacherConfig(str(fixture / 'teacher'), dtype=args.dtype, + autotp_size=args.tp, offload_to_cpu=False), + world_size=world, vocab_parallel=True) + layout, teacher.vocab_layout = prepare_vocab_parallel(candidate.module, teacher._callable.module, + tokenizer, tokenizer) + candidate._opsd_vocab_layout = layout + reference_layout = VocabShardLayout.from_model(dense.module) + layout.require_compatible(reference_layout) + handle = dense.module.get_output_embeddings().register_forward_hook( + lambda module, inputs, output: DenseGather.apply(output, reference_layout)) + cfg = OPSDConfig.from_dict({'student': {'model_name_or_path': 'random-tiny-Qwen2', 'dtype': args.dtype}, + 'teacher': {'model_name_or_path': str(fixture / 'teacher'), 'dtype': args.dtype, + 'autotp_size': args.tp, 'offload_to_cpu': False}, + 'distillation': {'loss_backend': 'vocab_parallel', 'loss_type': 'reverse_kl', + 'temperature': 0.5, 'chunk_size': 3}, + 'rollout': {'max_prompt_length': 16, 'max_response_length': 3}, + 'training': {'max_steps': 4 if args.phase == 'save' else 10, 'logging_steps': 1, 'save_steps': 1000}, + 'data': {'path': str(fixture / 'prompts.jsonl'), 'shuffle': False}}) + cfg.validate() + sampler = DistributedSampler(PromptDataset(cfg.data.path, tokenizer, 16), num_replicas=world // args.tp, + rank=dist.get_rank(candidate.data_parallel_group), shuffle=False) + dataset = sampler.dataset + loader = DataLoader(dataset, batch_size=1, sampler=sampler, drop_last=True, + collate_fn=LeftPaddedPromptCollator(tokenizer, 16)) + parameters = dict(candidate.module.named_parameters()) + references = dict(dense.module.named_parameters()) + assert parameters.keys() == references.keys() + initial = {name: master(p).clone() for name, p in parameters.items()} + arrays, records = {}, [] + identity = id(candidate.module.get_output_embeddings().weight) + + class AuditedTrainer(OPSDTrainer): + def _student_update(self, input_ids, attention_mask, response_mask, cache): + # A DP replica with no valid response positions must still join + # the TP KL collectives and Engine's GAS/DP reductions. + if args.empty_mask and self.step == 2 and dist.get_rank(candidate.data_parallel_group) == 1: + response_mask = torch.zeros_like(response_mask) + assert cache.cpu_logits.device.type == 'cpu' and cache.shape[-1] == layout.local_size + assert cache.cpu_logits.numel() == input_ids.numel() * layout.local_size + assert all(p.grad is None for p in teacher._callable.module.parameters()) + assert id(candidate.module.get_output_embeddings().weight) == identity + assert not getattr(candidate.module.get_output_embeddings(), '_opsd_rollout_gather_active', False) + dense.train() + logits = dense(input_ids=input_ids, attention_mask=attention_mask).logits[:, :-1].float() + target = gather(cache.cpu_logits.to(device=logits.device, dtype=dtype), layout)[:, :-1].float() + mask = response_mask[:, 1:].float() + tau = cfg.distillation.temperature + lp, lq = F.log_softmax(logits / tau, -1), F.log_softmax(target / tau, -1) + dense_loss = ((lp.exp() * (lp - lq)).sum(-1) * mask).sum() * tau**2 / mask.sum().clamp_min(1) + dense.backward(dense_loss) + gradient_errors = {} + original_step = candidate.step + + def checked_step(*step_args, **step_kwargs): + # ZeRO-2 exposes assembled gradients at the GAS boundary. + # Earlier microsteps retain internal accumulation fragments. + names = parameters if args.student_zero == 0 or candidate.is_gradient_accumulation_boundary() else [] + for name in names: + parameter = parameters[name] + actual, expected = safe_get_full_grad(parameter), safe_get_full_grad(references[name]) + assert actual is not None and expected is not None, name + gradient_errors[name] = compare(actual, expected, dtype, 'gradient ' + name) + arrays[f'micro{self.step}.grad.{name}'] = actual.detach().float().cpu().numpy().copy() + arrays[f'micro{self.step}.reference_grad.{name}'] = expected.detach().float().cpu().numpy().copy() + return original_step(*step_args, **step_kwargs) + + candidate.step = checked_step + try: + result = super()._student_update(input_ids, attention_mask, response_mask, cache) + finally: + candidate.step = original_step + dense.step() + loss_error = compare(result[0], dense_loss, dtype, 'trainer loss') + parameter_errors = {name: compare(master(p), master(references[name]), dtype, 'master ' + name) + for name, p in parameters.items()} + records.append({'microstep': self.step, 'global_steps': candidate.global_steps, + 'response_tokens': result[1], 'input_ids': input_ids.cpu().tolist(), + 'cache_shape': list(cache.shape), 'cache_bytes': cache.cpu_logits.numel() * 2, + 'loss_error': loss_error, 'max_gradient_error': max(gradient_errors.values(), default=0.0), + 'gradients_checked': len(gradient_errors), + 'max_master_parameter_error': max(parameter_errors.values()), + 'parameters_checked': len(parameters)}) + return result + + trainer = AuditedTrainer(cfg, candidate, teacher, tokenizer, + build_rollout(cfg.rollout, student_engine=candidate, tokenizer=tokenizer), loader) + checkpoint = Path(args.checkpoint) + if args.phase == 'resume': + for engine, name in [(candidate, 'candidate'), (dense, 'reference')]: + path, state = engine.load_checkpoint(str(checkpoint / name), tag='two-updates') + assert path and state['opsd_microstep'] == 4 and engine.global_steps == 2 + trainer.step = 4 + # Preserve the actual ordered DP data continuation across processes. + trainer.dataloader = DataLoader(Subset(dataset, list(iter(sampler))[4:]), batch_size=1, + collate_fn=LeftPaddedPromptCollator(tokenizer, 16)) + trainer.train() + if args.phase == 'save': + for engine, name in [(candidate, 'candidate'), (dense, 'reference')]: + engine.save_checkpoint(str(checkpoint / name), tag='two-updates', client_state={'opsd_microstep': 4}) + for name, parameter in parameters.items(): + arrays['final.' + name] = master(parameter).cpu().numpy() + arrays['reference_final.' + name] = master(references[name]).cpu().numpy() + arrays['delta.' + name] = (master(parameter) - initial[name]).cpu().numpy() + np.savez(out / ('rank-' + str(rank) + '.npz'), **arrays) + result = {'status': 'PASS', 'scope': 'Random tiny model actual trainer integration', 'rank': rank, + 'world': world, 'tp': args.tp, 'dp': world // args.tp, 'phase': args.phase, + 'optimizer_updates': candidate.global_steps, 'dtype': args.dtype, 'tied': args.tied, + 'student_zero_stage': args.student_zero, 'vocab_size': args.vocab_size, + 'parameters_checked': len(parameters), 'records': records, + 'teacher_zero_stage': teacher._callable.zero_optimization_stage(), + 'teacher_frozen': all(not p.requires_grad and p.grad is None for p in teacher._callable.module.parameters())} + (out / ('rank-' + str(rank) + '.json')).write_text(json.dumps(result, indent=2)) + handle.remove() + dist.barrier() + dist.destroy_process_group() + + +if __name__ == '__main__': + main() diff --git a/training/opsd/tests/distributed/run_vocab_parallel_probe.py b/training/opsd/tests/distributed/run_vocab_parallel_probe.py new file mode 100644 index 000000000..6a0f071da --- /dev/null +++ b/training/opsd/tests/distributed/run_vocab_parallel_probe.py @@ -0,0 +1,213 @@ +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +"""Real AutoTP model probe; run with torchrun on two GPUs. + +Random tiny Qwen2 weights qualify integration mechanics, not pretrained +distillation quality or representative training performance. +""" + +import argparse +import copy +import json +import os +from datetime import timedelta +from pathlib import Path +from types import SimpleNamespace + +import deepspeed +from deepspeed import comm as dist +import torch +import torch.nn.functional as F +from transformers import Qwen2Config, Qwen2ForCausalLM + +from deepspeed.runtime.rollout import RolloutRequest, SamplingConfig +from deepspeed.runtime.rollout.hybrid_engine_rollout import HybridEngineRollout +from deepspeed.sequence.vocab_parallel_ops import vocab_parallel_kl_div +from vocab_parallel import VocabShardLayout, gathered_rollout_logits + + +class GatheredReference(torch.autograd.Function): + """Dense oracle on the same TP layout; one replicated loss per TP rank.""" + + @staticmethod + def forward(ctx, local_logits, layout): + ctx.start, ctx.end = layout.start, layout.end + widths = [right - left for left, right in layout.ranges] + padded = F.pad(local_logits.contiguous(), (0, max(widths) - local_logits.shape[-1])) + shards = [torch.empty_like(padded) for _ in widths] + dist.all_gather(shards, padded, group=layout.group) + return torch.cat([value[..., :width] for value, width in zip(shards, widths)], -1) + + @staticmethod + def backward(ctx, gradient): + return gradient[..., ctx.start:ctx.end].contiguous(), None + + +def make_model(dtype, tied, width, layers, vocab_size=64): + config = Qwen2Config(vocab_size=vocab_size, hidden_size=width, intermediate_size=width * 2, + num_hidden_layers=layers, num_attention_heads=4, + num_key_value_heads=4, max_position_embeddings=64, + attention_dropout=0.0, tie_word_embeddings=tied, + pad_token_id=0, eos_token_id=1, bos_token_id=2) + config._attn_implementation = "eager" + return Qwen2ForCausalLM(config).to(dtype=dtype) + + +def engine_config(dtype, train): + config = {"train_micro_batch_size_per_gpu": 1, "gradient_accumulation_steps": 2 if train else 1, + "tensor_parallel": {"autotp_size": 2, "vocab_parallel_lm_head": True}, + "zero_optimization": {"stage": 0}, "gradient_clipping": 0.0, + "bf16": {"enabled": dtype == torch.bfloat16}} + if train: + config["optimizer"] = {"type": "SGD", "params": {"lr": 0.05}} + return config + + +def compare(actual, expected, dtype, what): + atol, rtol = (1e-5, 1e-4) if dtype == torch.float32 else (2e-3, 1e-2) + torch.testing.assert_close(actual.float(), expected.float(), atol=atol, rtol=rtol, + msg=lambda message: what + ": " + message) + return float((actual.float() - expected.float()).abs().max().item()) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--out", required=True) + parser.add_argument("--dtype", choices=["float32", "bfloat16"], default="float32") + parser.add_argument("--tied", action="store_true") + args = parser.parse_args() + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + deepspeed.init_distributed(dist_backend="nccl", auto_mpi_discovery=False, + timeout=timedelta(seconds=120)) + rank = dist.get_rank() + assert dist.get_world_size() == 2, "This probe requires TP2, DP1" + dtype = getattr(torch, args.dtype) + torch.manual_seed(1781) + student = make_model(dtype, args.tied, 32, 1) + teacher = make_model(dtype, args.tied, 64, 2) + prompts = torch.tensor([[2, 3, 4]], device="cuda") + mask = torch.ones_like(prompts) + + # Untied heads force the initial global maximum onto the second shard. + if not args.tied: + student.eval() + with torch.no_grad(): + hidden = student.model(prompts.cpu()).last_hidden_state[0, -1] + student.lm_head.weight.zero_() + student.lm_head.weight[48].copy_(hidden) + reference = copy.deepcopy(student).cuda() + teacher_reference = copy.deepcopy(teacher).cuda().eval() + for parameter in teacher.parameters(): + parameter.requires_grad_(False) + for parameter in teacher_reference.parameters(): + parameter.requires_grad_(False) + student_engine, *_ = deepspeed.initialize(model=student, model_parameters=student.parameters(), + config=engine_config(dtype, True)) + teacher_engine, *_ = deepspeed.initialize(model=teacher, config=engine_config(dtype, False)) + student_layout = VocabShardLayout.from_model(student_engine.module) + teacher_layout = VocabShardLayout.from_model(teacher_engine.module) + student_layout.require_compatible(teacher_layout) + reference_engine = None + if dtype == torch.bfloat16: + # BF16 TP GEMMs and reductions round differently from a full-width + # unsharded forward. Pair the real engines on the same layout and + # optimizer, while the reference loss still uses ordinary dense math. + reference_engine, *_ = deepspeed.initialize(model=reference, model_parameters=reference.parameters(), + config=engine_config(dtype, True)) + reference = reference_engine.module + reference_layout = VocabShardLayout.from_model(reference) + student_layout.require_compatible(reference_layout) + reference.get_output_embeddings().register_forward_hook( + lambda module, args, output: GatheredReference.apply(output, reference_layout)) + head = student_engine.module.get_output_embeddings() + parameter_id, storage_ptr = id(head.weight), head.weight.data_ptr() + if args.tied: + assert student_engine.module.get_input_embeddings().weight is head.weight + rollout = HybridEngineRollout(student_engine, SimpleNamespace(pad_token_id=0, eos_token_id=1)) + records = [] + + def generate_and_compare(tag): + student_engine.eval() + reference.eval() + with torch.no_grad(): + expected = reference.generate(prompts, attention_mask=mask, max_new_tokens=3, + eos_token_id=None, pad_token_id=0, do_sample=False) + with gathered_rollout_logits(student_engine.module, student_layout): + actual = rollout.generate(RolloutRequest(prompts, mask), + SamplingConfig(max_new_tokens=3, temperature=0.0)).input_ids + # Match the rollout's EOS trimming, rather than hiding EOS differences. + expected, _ = rollout._pad_after_eos(expected, response_start=prompts.shape[1], + eos_token_id=1, pad_token_id=0) + assert torch.equal(actual, expected), (tag, actual.tolist(), expected.tolist()) + if tag == "before" and not args.tied: + assert int(actual[0, prompts.shape[1]]) == 48 + assert id(head.weight) == parameter_id and head.weight.data_ptr() == storage_ptr + assert not getattr(head, "_opsd_rollout_gather_active", False) + sharded = student_engine.module(input_ids=actual, attention_mask=torch.ones_like(actual)).logits + assert sharded.shape[-1] == student_layout.local_size + records.append({"case": tag + "_global_rollout", "tokens": actual.tolist(), "status": "PASS"}) + return actual + + input_ids = generate_and_compare("before") + attention = torch.ones_like(input_ids) + response_mask = torch.zeros_like(input_ids[:, 1:], dtype=torch.float32) + response_mask[:, prompts.shape[1] - 1:] = 1 + optimizer = torch.optim.SGD(reference.parameters(), lr=0.05) + reference.train() + teacher_engine.eval() + for update in range(5): + for microstep in range(2): + student_engine.train() + output = student_engine(input_ids=input_ids, attention_mask=attention).logits[:, :-1].float() + with torch.no_grad(): + target = teacher_engine(input_ids=input_ids, attention_mask=attention).logits[:, :-1].float() + target_full = (GatheredReference.apply(target, teacher_layout) if reference_engine is not None else + teacher_reference(input_ids=input_ids, attention_mask=attention).logits[:, :-1].float()) + ref_logits = reference(input_ids=input_ids, attention_mask=attention).logits[:, :-1].float() + compare(output, ref_logits[..., student_layout.start:student_layout.end], dtype, "student shard") + compare(target, target_full[..., teacher_layout.start:teacher_layout.end], dtype, "teacher shard") + for reverse in (False, True): + candidate_tokens = vocab_parallel_kl_div(output, target, reverse=reverse, + tp_group=student_layout.group) + p, q = F.log_softmax(ref_logits, -1), F.log_softmax(target_full, -1) + expected_tokens = (p.exp() * (p - q)).sum(-1) if reverse else (q.exp() * (q - p)).sum(-1) + error = compare(candidate_tokens, expected_tokens, dtype, "KL direction") + records.append({"case": "KL", "update": update, "microstep": microstep, + "reverse": reverse, "max_error": error, "status": "PASS"}) + loss = (candidate_tokens * response_mask).sum() / response_mask.sum() + ref_loss = (expected_tokens * response_mask).sum() / response_mask.sum() + student_engine.backward(loss) + if reference_engine is None: + (ref_loss / 2).backward() + else: + reference_engine.backward(ref_loss) + student_engine.step() + if reference_engine is not None: + reference_engine.step() + if reference_engine is None: + optimizer.step() + optimizer.zero_grad() + reference_weight = reference.get_output_embeddings().weight + if reference_engine is None: + reference_weight = reference_weight[student_layout.start:student_layout.end] + compare(head.weight, reference_weight, + dtype, "optimizer head update") + assert all(parameter.grad is None for parameter in teacher_engine.module.parameters()) + records.append({"case": "optimizer_update", "update": update + 1, "status": "PASS"}) + generate_and_compare("after_five_updates") + result = {"status": "PASS", "scope": "Random tiny Qwen2 integration probe", "rank": rank, + "dtype": args.dtype, "tied": args.tied, "records": records, "optimizer_updates": 5, + "deepspeed_origin": deepspeed.__file__, "torch": torch.__version__, + "reference_layout": "same AutoTP layout with dense loss" if reference_engine is not None else "unsharded FP32", + "vocab_layout": {"vocab_size": student_layout.vocab_size, "start": student_layout.start, + "end": student_layout.end, "ranks": student_layout.ranks}} + out = Path(args.out) + out.mkdir(parents=True, exist_ok=True) + (out / f"rank-{rank}.json").write_text(json.dumps(result, indent=2)) + dist.barrier() + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/training/opsd/tests/test_vocab_parallel.py b/training/opsd/tests/test_vocab_parallel.py new file mode 100644 index 000000000..5f996b7a5 --- /dev/null +++ b/training/opsd/tests/test_vocab_parallel.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +"""Catch direction, temperature, mask, cache and lazy-dependency regressions.""" + +import builtins +from dataclasses import replace + +import pytest +import torch + +from config import OPSDConfig +from losses import streamed_distillation_loss +from teacher import TeacherLogitCache +from vocab_parallel import VocabShardLayout + + +def local_layout(width): + return VocabShardLayout(width, 0, width, (0,), ((0, width),), None) + + +@pytest.mark.parametrize("direction", ["forward_kl", "reverse_kl"]) +@pytest.mark.parametrize("temperature", [0.5, 1.0, 2.0]) +@pytest.mark.parametrize("chunk", [1, 3, 12]) +def test_streamed_local_kl_matches_independent_dense_oracle(direction, temperature, chunk): + torch.manual_seed(710) + student = torch.randn(2, 7, 11, requires_grad=True) + reference = student.detach().double().requires_grad_() + teacher = torch.randn(2, 7, 11) + mask = torch.tensor([[1, 0, 1, 1, 0, 1, 1], [0, 1, 1, 0, 1, 0, 1]], dtype=torch.float32) + layout = local_layout(11) + cache = TeacherLogitCache.from_gpu_logits(teacher, store_dtype=torch.bfloat16, layout=layout) + quantized = teacher.to(torch.bfloat16).double() + logp = torch.log_softmax(reference / temperature, dim=-1) + logq = torch.log_softmax(quantized / temperature, dim=-1) + tokens = (logq.exp() * (logq - logp)).sum(-1) if direction == "forward_kl" else (logp.exp() * (logp - logq)).sum(-1) + expected = (tokens * mask).sum() * temperature**2 / mask.sum() + result = streamed_distillation_loss(student, lambda a, b: cache.chunk_to_device(a, b, student.device), + mask, loss_type=direction, temperature=temperature, chunk_size=chunk, + loss_backend="vocab_parallel", vocab_layout=layout) + result.backward() + expected.backward() + torch.testing.assert_close(result.double(), expected, atol=1e-6, rtol=1e-5) + torch.testing.assert_close(student.grad.double(), reference.grad, atol=1e-6, rtol=1e-5) + + +def test_zero_mask_is_backward_safe_and_does_not_skip_collective_chunks(): + student = torch.randn(1, 4, 9, requires_grad=True) + fetched = [] + + def fetch(start, end): + fetched.append((start, end)) + return torch.zeros(1, end - start, 9) + + result = streamed_distillation_loss(student, fetch, torch.zeros(1, 4), chunk_size=3, + loss_backend="vocab_parallel", vocab_layout=local_layout(9)) + result.backward() + assert result.item() == 0 + assert fetched == [(0, 3), (3, 4)] + assert torch.count_nonzero(student.grad) == 0 + + +def test_cache_binds_actual_width_and_rejects_use_after_free(): + with pytest.raises(ValueError, match="width"): + TeacherLogitCache(torch.zeros(1, 3, 8), layout=local_layout(9)) + cache = TeacherLogitCache(torch.zeros(1, 3, 9), layout=local_layout(9)) + cache.free() + with pytest.raises(RuntimeError, match="freed"): + cache.chunk_to_device(0, 1, torch.device("cpu")) + + +def test_dense_does_not_import_optional_kl_dependency(monkeypatch): + original = builtins.__import__ + + def blocked(name, *args, **kwargs): + if name == "deepspeed.sequence.vocab_parallel_ops": + raise ImportError("dependency absent") + return original(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", blocked) + logits = torch.randn(1, 2, 5, requires_grad=True) + fetch = lambda a, b: torch.zeros(1, b - a, 5) + assert torch.isfinite(streamed_distillation_loss(logits, fetch, torch.ones(1, 2))) + with pytest.raises(RuntimeError, match="PR #8791"): + streamed_distillation_loss(logits, fetch, torch.ones(1, 2), loss_backend="vocab_parallel", + vocab_layout=local_layout(5)) + + +def test_old_config_defaults_dense_and_opt_in_rejects_unsupported_routes(): + cfg = OPSDConfig.from_dict({"student": {"model_name_or_path": "student"}, + "teacher": {"model_name_or_path": "teacher"}}) + assert cfg.distillation.loss_backend == "dense" + cfg.validate() + cfg.distillation = replace(cfg.distillation, loss_backend="vocab_parallel", loss_type="jsd") + with pytest.raises(ValueError, match="JSD"): + cfg.validate() + cfg.distillation.loss_type = "reverse_kl" + cfg.rollout.temperature = 1.0 + with pytest.raises(ValueError, match="greedy"): + cfg.validate() + cfg.rollout.temperature = 0.0 + cfg.validate() + for feature in ["use_graph_capture", "use_shared_prefill"]: + setattr(cfg.rollout, feature, True) + with pytest.raises(ValueError, match="eager"): + cfg.validate() + setattr(cfg.rollout, feature, False) diff --git a/training/opsd/trainer.py b/training/opsd/trainer.py index 173f0d2a1..a77a53a7a 100644 --- a/training/opsd/trainer.py +++ b/training/opsd/trainer.py @@ -25,6 +25,7 @@ import os import time from abc import ABC, abstractmethod +from contextlib import nullcontext from typing import Any import torch @@ -34,6 +35,7 @@ from config import OPSDConfig from losses import streamed_distillation_loss from utils import build_response_mask +from vocab_parallel import gathered_rollout_logits, validate_tp_batch, validate_training_logits from deepspeed.runtime.rollout import RolloutEngine, RolloutRequest, SamplingConfig @@ -73,6 +75,7 @@ def __init__( self.device = get_accelerator().current_device_name() self.step = 0 + self.vocab_layout = getattr(student_engine, "_opsd_vocab_layout", None) # ------------------------------------------------------------------ # Driver @@ -117,14 +120,19 @@ def _train_step(self, batch) -> dict: top_k=self.cfg.rollout.top_k, n_samples_per_prompt=self.cfg.rollout.n_samples_per_prompt, ) - roll = self.rollout.generate( - RolloutRequest(prompt_ids=prompt_ids, prompt_attention_mask=prompt_attn), - sampling, - ) + scope = (gathered_rollout_logits(self.student_engine.module, self.vocab_layout) + if self.vocab_layout is not None else nullcontext()) + with scope: + roll = self.rollout.generate( + RolloutRequest(prompt_ids=prompt_ids, prompt_attention_mask=prompt_attn), + sampling, + ) input_ids = roll.input_ids.to(self.device, non_blocking=True) attention_mask = roll.attention_mask.to(self.device, non_blocking=True) response_start_idx = roll.response_start_idx.to(self.device, non_blocking=True) response_mask = build_response_mask(response_start_idx, attention_mask) + if self.vocab_layout is not None: + validate_tp_batch(self.vocab_layout, input_ids, attention_mask, response_mask) t_rollout = time.time() - t_start # --- Phase 1: teacher forward → host-cached logits ---------------- @@ -132,11 +140,35 @@ def _train_step(self, batch) -> dict: teacher_cache = self.teacher.forward_to_cache(input_ids, attention_mask) t_teacher = time.time() - t1 + try: + loss, response_tokens, t_student = self._student_update(input_ids, attention_mask, response_mask, + teacher_cache) + finally: + teacher_cache.free() + + loss_for_log = loss.detach().clone() + if dist.is_initialized(): + dist.all_reduce(loss_for_log) + loss_for_log /= dist.get_world_size() + + return { + "loss": float(loss_for_log.item()), + "rollout_s": t_rollout, + "teacher_s": t_teacher, + "student_s": t_student, + "step_s": time.time() - t_start, + "response_tokens": response_tokens, + } + + def _student_update(self, input_ids, attention_mask, response_mask, teacher_cache): + # --- Phase 2: student forward + streamed KL + backward ------------ t2 = time.time() self.student_engine.train() outputs = self.student_engine(input_ids=input_ids, attention_mask=attention_mask) student_logits = outputs.logits # [B, T, V] + if self.vocab_layout is not None: + validate_training_logits(self.vocab_layout, student_logits, teacher_cache) # Shift for next-token prediction: logits at position t predict token # at t+1, so the loss aligns student_logits[:, :-1] with the position @@ -161,28 +193,15 @@ def _fetch(start: int, end: int) -> torch.Tensor: loss_type=self.cfg.distillation.loss_type, temperature=self.cfg.distillation.temperature, chunk_size=self.cfg.distillation.chunk_size, + loss_backend=self.cfg.distillation.loss_backend, + vocab_layout=self.vocab_layout, ) self.student_engine.backward(loss) self.student_engine.step() - teacher_cache.free() t_student = time.time() - t2 - - # Reduce loss across ranks for a clean log line. - loss_for_log = loss.detach().clone() - if dist.is_initialized(): - dist.all_reduce(loss_for_log) - loss_for_log /= dist.get_world_size() - - return { - "loss": float(loss_for_log.item()), - "rollout_s": t_rollout, - "teacher_s": t_teacher, - "student_s": t_student, - "step_s": time.time() - t_start, - "response_tokens": int(mask_shifted.sum().item()), - } + return loss, int(mask_shifted.sum().item()), t_student # ------------------------------------------------------------------ # Logging / checkpointing diff --git a/training/opsd/vocab_parallel.py b/training/opsd/vocab_parallel.py new file mode 100644 index 000000000..acef77236 --- /dev/null +++ b/training/opsd/vocab_parallel.py @@ -0,0 +1,179 @@ +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +"""Vocabulary layout and eager rollout support for sharded OPSD logits.""" + +from contextlib import contextmanager +from dataclasses import dataclass, field +import hashlib +import json +from typing import Any + +import torch + + +@dataclass(frozen=True) +class VocabShardLayout: + vocab_size: int + start: int + end: int + ranks: tuple[int, ...] + ranges: tuple[tuple[int, int], ...] + group: Any = field(compare=False, repr=False) + + @property + def local_size(self) -> int: + return self.end - self.start + + @classmethod + def from_model(cls, model) -> "VocabShardLayout": + from deepspeed import comm as dist + + error = None + try: + head = model.get_output_embeddings() + if head is None: + raise ValueError("The model must expose its output embedding") + sharded = getattr(head, "is_vocab_parallel_lm_head", False) + group = getattr(head, "mp_group", None) if sharded else None + size = dist.get_world_size(group) if group is not None else 1 + weight_shape = getattr(head.weight, "ds_shape", head.weight.shape) + vocab = int(head.vocab_size if sharded else weight_shape[0]) + start = int(head.vocab_start_index) if sharded else 0 + end = int(head.vocab_end_index) if sharded else vocab + width = int(weight_shape[0]) + except (ValueError, AttributeError, IndexError, TypeError) as failure: + error = str(failure) + # Check local metadata before TP peers enter the tensor collective. + # A malformed head on one rank must produce a uniform startup error. + _raise_if_any_error(error, "Vocabulary head metadata") + rank = dist.get_rank() if dist.is_initialized() else 0 + local = torch.tensor([rank, vocab, start, end, width], + device=head.weight.device, dtype=torch.int64) + gathered = [torch.empty_like(local) for _ in range(size)] + if size > 1: + dist.all_gather(gathered, local, group=group) + else: + gathered[0].copy_(local) + metadata = [entry.tolist() for entry in gathered] + cursor, error = 0, None + for _, total, left, right, width in metadata: + if total != vocab or left != cursor or right <= left or right - left != width: + error = "Vocabulary shards must form a nonempty, contiguous partition" + cursor = right + if cursor != vocab: + error = "Vocabulary shards do not cover the model vocabulary" + _raise_if_any_error(error, "Vocabulary partition") + return cls(vocab, start, end, tuple(m[0] for m in metadata), + tuple((m[2], m[3]) for m in metadata), group) + + def require_compatible(self, other: "VocabShardLayout") -> None: + if (self.vocab_size, self.ranks, self.ranges, self.start, self.end) != ( + other.vocab_size, other.ranks, other.ranges, other.start, other.end): + raise ValueError("Teacher and student require the same ordered TP vocabulary partition") + + +def _agree_metadata(entry): + from deepspeed import comm as dist + + if not dist.is_initialized(): + return [entry] + entries = [None] * dist.get_world_size() + dist.all_gather_object(entries, entry) + return entries + + +def _raise_if_any_error(error, context): + errors = [item["error"] for item in _agree_metadata({"error": error}) if item["error"]] + if errors: + raise ValueError(context + " failed: " + "; ".join(sorted(set(errors)))) + + +def tokenizer_fingerprint(tokenizer): + semantic = {"vocab": tokenizer.get_vocab(), "added": tokenizer.get_added_vocab(), + "special_ids": {name: getattr(tokenizer, name, None) + for name in ["pad_token_id", "eos_token_id", "bos_token_id", "unk_token_id"]}} + return hashlib.sha256(json.dumps(semantic, sort_keys=True).encode()).hexdigest() + + +def prepare_vocab_parallel(student, teacher, student_tokenizer, teacher_tokenizer): + student_layout = VocabShardLayout.from_model(student) + teacher_layout = VocabShardLayout.from_model(teacher) + error = None + try: + student_layout.require_compatible(teacher_layout) + if tokenizer_fingerprint(student_tokenizer) != tokenizer_fingerprint(teacher_tokenizer): + raise ValueError("Teacher and student token IDs have different vocabulary semantics") + except ValueError as failure: + error = str(failure) + errors = [item["error"] for item in _agree_metadata({"error": error}) if item["error"]] + if errors: + raise ValueError("Vocabulary-parallel preflight failed: " + "; ".join(sorted(set(errors)))) + return student_layout, teacher_layout + + +def validate_tp_batch(layout, input_ids, attention_mask, response_mask): + from deepspeed import comm as dist + + rank = dist.get_rank() if dist.is_initialized() else 0 + digests = [hashlib.sha256(value.detach().cpu().contiguous().numpy().tobytes()).hexdigest() + for value in [input_ids, attention_mask, response_mask]] + entry = {"rank": rank, "ranks": layout.ranks, "shape": list(input_ids.shape), "digests": digests} + entries = _agree_metadata(entry) + by_rank = {item["rank"]: item for item in entries} + for item in entries: + for peer in item["ranks"]: + if (by_rank[peer]["shape"], by_rank[peer]["digests"]) != (item["shape"], item["digests"]): + raise ValueError("Input tokens and masks must agree within each TP group") + + +def validate_training_logits(layout, student_logits, teacher_cache): + error = None + try: + layout.require_compatible(teacher_cache.layout) + if student_logits.ndim != 3 or tuple(student_logits.shape) != teacher_cache.shape: + raise ValueError("Teacher cache and student logits require identical local shapes") + if student_logits.shape[-1] != layout.local_size: + raise ValueError("Model output differs from the validated local vocabulary width") + except (ValueError, AttributeError) as failure: + error = str(failure) + errors = [item["error"] for item in _agree_metadata({"error": error}) if item["error"]] + if errors: + raise ValueError("Vocabulary-parallel forward failed: " + "; ".join(sorted(set(errors)))) + + +@contextmanager +def gathered_rollout_logits(model, layout: VocabShardLayout): + """Expose only the last token's full vocabulary during eager generation. + + Training and teacher forwards keep their local output shards. The hook + neither replaces parameters nor changes the head's gather configuration. + """ + from deepspeed import comm as dist + + head = model.get_output_embeddings() + if getattr(head, "_opsd_rollout_gather_active", False): + raise RuntimeError("Vocabulary rollout scopes cannot be nested") + + def gather_last_token(module, args, output): + if model.training or torch.is_grad_enabled(): + raise RuntimeError("Vocabulary rollout gathering requires eval and no_grad") + if output.ndim != 3 or output.shape[-1] != layout.local_size: + raise ValueError("Rollout head output does not match its vocabulary layout") + last = output[:, -1:, :].contiguous() + if len(layout.ranks) == 1: + return last + widths = [right - left for left, right in layout.ranges] + padded = torch.nn.functional.pad(last, (0, max(widths) - last.shape[-1])) + shards = [torch.empty_like(padded) for _ in widths] + dist.all_gather(shards, padded, group=layout.group) + return torch.cat([shard[..., :width] for shard, width in zip(shards, widths)], dim=-1) + + head._opsd_rollout_gather_active = True + handle = None + try: + handle = head.register_forward_hook(gather_last_token) + yield + finally: + if handle is not None: + handle.remove() + del head._opsd_rollout_gather_active From 94fe46bcf4109f2377622970603b3099cfb64ad7 Mon Sep 17 00:00:00 2001 From: 0z5a Date: Wed, 14 Oct 2026 20:30:00 +0200 Subject: [PATCH 2/2] Validate OPSD benchmark labels and prepare result paths Signed-off-by: 0z5a --- training/opsd/README.md | 7 +++- .../benchmarks/bench_vocab_parallel_opsd.py | 24 +++++++++++ training/opsd/tests/test_vocab_parallel.py | 41 +++++++++++++++++++ 3 files changed, 71 insertions(+), 1 deletion(-) diff --git a/training/opsd/README.md b/training/opsd/README.md index 79a403229..3082956e1 100644 --- a/training/opsd/README.md +++ b/training/opsd/README.md @@ -200,7 +200,12 @@ the harness explicitly selects the dependency's documented full-logits configuration for both producers, including the upstream teacher initializer. `gathered_reference` uses the same candidate head layout with dense math and a local-slice gather backward; `vocab_parallel` retains local logits. Replay -excludes generation. Online always generates new responses. Report actual +excludes generation. Online always generates new responses. The benchmark +checks that `--implementation` matches the configured loss backend before +initializing models: `upstream_dense`/`dense_default` require `dense` (or the +omitted default), while `gathered_reference`/`vocab_parallel` require +`vocab_parallel` to construct the sharded heads. It also creates and checks +the result directory before starting GPU work. Report actual optimizer updates, unique response tokens counted once per TP group, each run's workload hash, cache payload and device peak; distinguish these from the original trainer's unsynchronized diagnostic phase timers. diff --git a/training/opsd/benchmarks/bench_vocab_parallel_opsd.py b/training/opsd/benchmarks/bench_vocab_parallel_opsd.py index aac4bb8eb..5846ff99a 100644 --- a/training/opsd/benchmarks/bench_vocab_parallel_opsd.py +++ b/training/opsd/benchmarks/bench_vocab_parallel_opsd.py @@ -41,6 +41,26 @@ def backward(ctx, gradient): return gradient[..., ctx.start:ctx.end].contiguous(), None +def prepare_benchmark_inputs(implementation, config_path, output_jsonl): + with Path(config_path).open() as source: + config = json.load(source) + sharded = implementation in ('vocab_parallel', 'gathered_reference') + expected_backend = 'vocab_parallel' if sharded else 'dense' + backend = config.get('distillation', {}).get('loss_backend', 'dense') + if backend != expected_backend: + raise ValueError(f'{implementation} requires distillation.loss_backend={expected_backend!r}; got {backend!r}') + with Path(config['deepspeed_config']).open() as source: + ds_config = json.load(source) + if (implementation == 'dense_default' + and ds_config.get('tensor_parallel', {}).get('vocab_parallel_lm_head') is True): + raise ValueError('dense_default requires full-vocabulary heads; disable vocab_parallel_lm_head') + output = Path(output_jsonl) + output.parent.mkdir(parents=True, exist_ok=True) + # Fail before model initialization if the result path cannot be written. + with output.open('a'): + pass + + def run(): parser = argparse.ArgumentParser() parser.add_argument('--mode', choices=['online', 'replay'], required=True) @@ -57,6 +77,10 @@ def run(): help='Use the same PyTorch SDPA backend for every measured implementation') parser.add_argument('--profile', action='store_true', help='Additional phase synchronization; separate from throughput') args = parser.parse_args() + try: + prepare_benchmark_inputs(args.implementation, args.config, args.output_jsonl) + except (ValueError, OSError) as error: + parser.error(str(error)) if args.implementation == 'upstream_dense': if not args.upstream_opsd_dir: parser.error('upstream_dense needs the unchanged frozen upstream OPSD directory') diff --git a/training/opsd/tests/test_vocab_parallel.py b/training/opsd/tests/test_vocab_parallel.py index 5f996b7a5..ee68b7139 100644 --- a/training/opsd/tests/test_vocab_parallel.py +++ b/training/opsd/tests/test_vocab_parallel.py @@ -104,3 +104,44 @@ def test_old_config_defaults_dense_and_opt_in_rejects_unsupported_routes(): with pytest.raises(ValueError, match="eager"): cfg.validate() setattr(cfg.rollout, feature, False) + + +@pytest.mark.parametrize("implementation,backend", [ + ("upstream_dense", "dense"), ("dense_default", "dense"), + ("gathered_reference", "vocab_parallel"), ("vocab_parallel", "vocab_parallel"), +]) +def test_benchmark_requires_matching_backend_and_creates_output(tmp_path, implementation, backend): + import json + from benchmarks.bench_vocab_parallel_opsd import prepare_benchmark_inputs + + ds_path = tmp_path / "ds.json" + ds_path.write_text(json.dumps({"tensor_parallel": {"autotp_size": 2}})) + config_path = tmp_path / "opsd.json" + config = {"deepspeed_config": str(ds_path), "distillation": {"loss_backend": backend}} + config_path.write_text(json.dumps(config)) + output = tmp_path / "fresh" / "nested" / "results.jsonl" + prepare_benchmark_inputs(implementation, config_path, output) + assert output.is_file() + output.write_text("existing result\n") + prepare_benchmark_inputs(implementation, config_path, output) + assert output.read_text() == "existing result\n" + + config["distillation"]["loss_backend"] = "dense" if backend == "vocab_parallel" else "vocab_parallel" + config_path.write_text(json.dumps(config)) + with pytest.raises(ValueError, match="requires distillation.loss_backend"): + prepare_benchmark_inputs(implementation, config_path, output) + + +def test_benchmark_rejects_sharded_dense_heads_and_unwritable_output(tmp_path): + import json + from benchmarks.bench_vocab_parallel_opsd import prepare_benchmark_inputs + + ds_path = tmp_path / "ds.json" + ds_path.write_text(json.dumps({"tensor_parallel": {"autotp_size": 2, "vocab_parallel_lm_head": True}})) + config_path = tmp_path / "opsd.json" + config_path.write_text(json.dumps({"deepspeed_config": str(ds_path)})) + with pytest.raises(ValueError, match="full-vocabulary"): + prepare_benchmark_inputs("dense_default", config_path, tmp_path / "results.jsonl") + ds_path.write_text("{}") + with pytest.raises(IsADirectoryError): + prepare_benchmark_inputs("dense_default", config_path, tmp_path)