From f74f9906b8237b29a67ce253e594ad3e435c4240 Mon Sep 17 00:00:00 2001 From: 0z5a Date: Wed, 14 Oct 2026 19:44:18 +0200 Subject: [PATCH 1/4] 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/4] 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) From e9fd43ca0125264da0bb929c7783e9f5bf63e9b4 Mon Sep 17 00:00:00 2001 From: 0z5a Date: Sat, 17 Oct 2026 17:07:47 +0200 Subject: [PATCH 3/4] Skip OPSD chunks without response tokens on the vocab-parallel route The vocabulary-parallel KL route fetched and computed every sequence chunk, including prompt-only chunks whose loss is masked to zero, because skipping them on a per-rank decision could leave the op's collectives unmatched. With a 2048-token prompt and at most 64 response tokens that is 9 chunks fetched per microstep where 1-2 are needed, which made this route slower than gathering the full logits. The TP group now agrees on the chunks with response tokens through one MAX all-reduce of a per-chunk flag, then skips the rest; the dense route skips the same chunks as before with one host sync instead of one per chunk. The shards are also passed to the op in bf16 (the student chunk made contiguous) instead of as fp32 copies, since the op computes in fp32 internally; the copies doubled what its kernels read and what autograd keeps per chunk. Only a temperature still needs the fp32 division first. Co-Authored-By: Claude Opus 5.5 Signed-off-by: 0z5a --- training/opsd/losses.py | 27 ++++++++++++++++------ training/opsd/tests/test_vocab_parallel.py | 11 +++++++-- 2 files changed, 29 insertions(+), 9 deletions(-) diff --git a/training/opsd/losses.py b/training/opsd/losses.py index 033f7aad9..e3b7ce651 100644 --- a/training/opsd/losses.py +++ b/training/opsd/losses.py @@ -176,11 +176,21 @@ def streamed_distillation_loss( total_tokens = mask_f.sum().clamp_min(1.0) total_loss = student_logits[:, :0].sum(dtype=torch.float32) - for start in range(0, T, chunk_size): + # Chunks without response tokens contribute nothing, so they are neither fetched nor computed. + # The vocabulary-parallel op runs collectives, so its TP group agrees on the chunks to skip; that + # stays symmetric even if a caller's masks differ across TP ranks. + starts = list(range(0, T, chunk_size)) + tokens_per_position = F.pad(mask_f.sum(dim=0), (0, len(starts) * chunk_size - T)) + has_tokens = (tokens_per_position.view(len(starts), chunk_size).sum(dim=1) > 0).to(torch.int32) + if parallel_kl is not None and len(vocab_layout.ranks) > 1: + from deepspeed import comm as dist + dist.all_reduce(has_tokens, op=dist.ReduceOp.MAX, group=vocab_layout.group) + + for start, needed in zip(starts, has_tokens.tolist()): + if not needed: + continue end = min(start + chunk_size, T) chunk_mask = mask_f[:, start:end] - if parallel_kl is None and chunk_mask.sum().item() == 0: - continue teacher_chunk = teacher_chunk_fetcher(start, end) if teacher_chunk.shape != student_logits[:, start:end].shape: raise RuntimeError(f"fetcher returned chunk of length {teacher_chunk.shape[1]}, " @@ -188,11 +198,14 @@ def streamed_distillation_loss( 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() + # The op computes in fp32 internally; fp32 copies would double what its kernels read and what + # autograd keeps per chunk. Only a temperature needs the fp32 division first. + student_chunk = student_logits[:, start:end] + teacher_chunk = teacher_chunk.detach() if temperature != 1.0: - student_chunk = student_chunk / temperature - teacher_chunk = teacher_chunk / temperature + student_chunk = student_chunk.float() / temperature + teacher_chunk = teacher_chunk.float() / temperature + student_chunk = student_chunk.contiguous() per_tok = parallel_kl(student_chunk, teacher_chunk, reverse=loss_type == "reverse_kl", tp_group=vocab_layout.group) per_tok = per_tok * temperature**2 diff --git a/training/opsd/tests/test_vocab_parallel.py b/training/opsd/tests/test_vocab_parallel.py index ee68b7139..7651362f4 100644 --- a/training/opsd/tests/test_vocab_parallel.py +++ b/training/opsd/tests/test_vocab_parallel.py @@ -43,7 +43,7 @@ def test_streamed_local_kl_matches_independent_dense_oracle(direction, temperatu 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(): +def test_chunks_without_response_tokens_are_skipped_and_backward_safe(): student = torch.randn(1, 4, 9, requires_grad=True) fetched = [] @@ -55,9 +55,16 @@ def fetch(start, end): loss_backend="vocab_parallel", vocab_layout=local_layout(9)) result.backward() assert result.item() == 0 - assert fetched == [(0, 3), (3, 4)] + assert fetched == [] assert torch.count_nonzero(student.grad) == 0 + student.grad = None + result = streamed_distillation_loss(student, fetch, torch.tensor([[0.0, 0.0, 0.0, 1.0]]), chunk_size=3, + loss_backend="vocab_parallel", vocab_layout=local_layout(9)) + result.backward() + assert fetched == [(3, 4)] + assert torch.count_nonzero(student.grad[:, :3]) == 0 + def test_cache_binds_actual_width_and_rejects_use_after_free(): with pytest.raises(ValueError, match="width"): From cea3bf7a61d4eeca07c1f3032e954cd5e043da4e Mon Sep 17 00:00:00 2001 From: 0z5a Date: Sat, 17 Oct 2026 17:36:49 +0200 Subject: [PATCH 4/4] Copy OPSD teacher cache chunks to the device from pinned memory The teacher cache is a pinned [B, T, V] host tensor, and each loss chunk was copied to the device with cpu_logits[:, start:end].to(device). For B > 1 that slice is strided. PyTorch copies a strided host tensor to a GPU through a temporary: it first makes the slice contiguous in newly allocated pageable host memory, then copies the temporary synchronously. So every chunk paid a host copy into fresh pages plus a pageable transfer, even though the cache itself is pinned. Each row's slice cpu_logits[b, start:end] is contiguous within the pinned buffer, so the chunk is now copied to a contiguous device tensor one row at a time. When the cache is pinned the row copies are non-blocking. The pinned host allocator records the copies on the slices' storage, so freeing the cache with copies still queued keeps the buffer until they finish. A dtype override is applied on the device after the copy. Host destinations keep the previous slice-and-convert path. This needs no staging buffer and no change to the cache layout, its lifetime or free(). The bytes are unchanged, so losses and updates are bitwise identical. A CUDA test checks the device chunks against the host slices bit for bit: one and three rows, pinned and pageable caches, the last partial chunk, and both dtype overrides. Co-Authored-By: Claude Opus 5.5 Signed-off-by: 0z5a --- training/opsd/teacher.py | 18 ++++++++++++++---- training/opsd/tests/test_teacher_caching.py | 21 +++++++++++++++++++++ 2 files changed, 35 insertions(+), 4 deletions(-) diff --git a/training/opsd/teacher.py b/training/opsd/teacher.py index 9afb29ec0..3a93fe6b8 100644 --- a/training/opsd/teacher.py +++ b/training/opsd/teacher.py @@ -91,14 +91,24 @@ def chunk_to_device(self, """Slice ``[:, start:end, :]`` and stage it on ``device``. ``dtype`` is the dtype on the destination; if ``None``, the stored - dtype is preserved. + dtype is preserved. Off-host destinations convert after the copy. """ - _, T, _ = self.shape + B, T, _ = self.shape 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=False) - return out + dtype = dtype if dtype is not None else chunk.dtype + if torch.device(device).type == "cpu": + return chunk.to(dtype=dtype) + # For B > 1 the slice is strided, and copying it to a device first makes a + # contiguous temporary in pageable memory. Each row's slice is contiguous + # in the pinned buffer, so copy rows directly. The pinned allocator keeps + # the buffer until the copies finish, even if the cache is freed earlier. + pinned = self.cpu_logits.is_pinned() + out = torch.empty(chunk.shape, dtype=chunk.dtype, device=device) + for row in range(B): + out[row].copy_(chunk[row], non_blocking=pinned) + return out.to(dtype=dtype) def free(self) -> None: """Drop the underlying buffer so a step's teacher cache can be GC'd diff --git a/training/opsd/tests/test_teacher_caching.py b/training/opsd/tests/test_teacher_caching.py index 04aebfad5..078b5cf11 100644 --- a/training/opsd/tests/test_teacher_caching.py +++ b/training/opsd/tests/test_teacher_caching.py @@ -58,6 +58,27 @@ def test_rejects_gpu_resident_logits(): TeacherLogitCache(cpu_logits=torch.zeros(1, 8, 4, device="cuda")) +def test_device_chunks_match_host_slices(): + if not torch.cuda.is_available(): #ignore-cuda + pytest.skip("no CUDA available to stage chunks on a GPU") + torch.manual_seed(0) + for batch in [1, 3]: + src = torch.randn(batch, 50, 24) + for store, override in [(torch.bfloat16, torch.float32), (torch.float32, torch.bfloat16)]: + pinned = TeacherLogitCache.from_gpu_logits(src, store_dtype=store) + pageable = TeacherLogitCache(cpu_logits=pinned.cpu_logits.clone()) + assert pinned.cpu_logits.is_pinned() and not pageable.cpu_logits.is_pinned() + for cache in [pinned, pageable]: + # Includes the last partial chunk and a single position. + for start, end in [(0, 16), (16, 32), (48, 50), (5, 6)]: + for dtype in [None, override]: + expected = cache.cpu_logits[:, start:end].to(dtype=dtype or store).contiguous() + got = cache.chunk_to_device(start, end, torch.device("cuda"), dtype=dtype) + assert got.is_contiguous() and got.dtype == expected.dtype + bits = torch.int16 if expected.dtype == torch.bfloat16 else torch.int32 + assert torch.equal(got.cpu().view(bits), expected.view(bits)) + + def test_dtype_override_in_chunk_to_device(): src = torch.randn(2, 8, 16, dtype=torch.float32) cache = TeacherLogitCache.from_gpu_logits(src, store_dtype=torch.float32)