From 8c483c781b51dbd44289b0b48e8a3133f0c97186 Mon Sep 17 00:00:00 2001 From: Ofir Ben Shoham Date: Tue, 30 Jun 2026 13:48:10 +0300 Subject: [PATCH 1/2] Add online target mode for DSpark training (no precomputed cache) Adds an opt-in online training path that recomputes target hidden states every step instead of reading them from a precomputed target cache. - Lift run_target_forward_with_hooks + backbone helpers out of scripts/data/prepare_target_cache.py into deepspec/modeling/target_extract.py so cache generation and online training share one extraction implementation. - base_trainer: when data.online_target is set, keep the frozen target backbone resident on-device, stream conversations via JsonLineDataset + ConversationCollator, and skip cache validation. - dspark_trainer.run_batch: recompute target hidden states via the shared helper when online; offline (cached) path is unchanged. - cuda_prefetcher: skip None micro-batches (ConversationCollator drops batches below min_loss_tokens) so downstream never sees None. - eagle3_trainer: assert online_target is unsupported (different layer-offset semantics; not yet implemented). - Add config/dspark/dspark_qwen3_4b_online.py example and a CPU equivalence test asserting online extraction matches the cache-path tensors, including a padding-invariance check. Offline cached training is fully preserved (online_target defaults to False). Note: GPU/end-to-end training was not run (no CUDA in dev env); validation is the CPU correctness test in tests/test_online_target_extract.py. --- config/dspark/dspark_qwen3_4b_online.py | 76 +++++++++++ deepspec/data/__init__.py | 2 + deepspec/data/cuda_prefetcher.py | 21 ++- deepspec/modeling/target_extract.py | 100 ++++++++++++++ deepspec/trainer/base_trainer.py | 57 ++++++-- deepspec/trainer/dspark_trainer.py | 17 +++ deepspec/trainer/eagle3_trainer.py | 7 + scripts/data/prepare_target_cache.py | 95 +------------- tests/test_online_target_extract.py | 165 ++++++++++++++++++++++++ 9 files changed, 433 insertions(+), 107 deletions(-) create mode 100644 config/dspark/dspark_qwen3_4b_online.py create mode 100644 deepspec/modeling/target_extract.py create mode 100644 tests/test_online_target_extract.py diff --git a/config/dspark/dspark_qwen3_4b_online.py b/config/dspark/dspark_qwen3_4b_online.py new file mode 100644 index 00000000..58fe2c0f --- /dev/null +++ b/config/dspark/dspark_qwen3_4b_online.py @@ -0,0 +1,76 @@ +import os +from deepspec.trainer import Qwen3DSparkTrainer +BASE_TB_DIR = os.path.expanduser("~/tensorboard") +BASE_CKPT_DIR = os.path.expanduser("~/checkpoints") +project_name = "deepspec" +exp_name = "dspark_block7_qwen3_4b_online" +seed = 42 + +model = dict( + target_model_name_or_path="Qwen/Qwen3-4B", + block_size=7, + num_draft_layers=5, + target_layer_ids=[1, 9, 17, 25, 33], + mask_token_id=151669, + num_anchors=512, + + ## markov head + markov_rank=256, + markov_head_type='vanilla', + + ## confidence head + confidence_head_alpha=1.0, + confidence_head_with_markov=True, + + ## loss + loss_decay_gamma=4.0, + ce_loss_alpha=0.1, + l1_loss_alpha=0.9, +) + +train = dict( + trainer_cls=Qwen3DSparkTrainer, + lr=6.0e-4, + warmup_ratio=0.04, + weight_decay=0.0, + precision="bf16", + local_batch_size=1, + global_batch_size=512, + num_train_epochs=10, + max_train_steps=None, + max_grad_norm=1.0, + sharding_strategy="no_shard", + torch_compile=True, +) + +logging = dict( + logging_steps=10, + checkpointing_steps=3000, +) + +data = dict( + # Online target mode: recompute target hidden states every step instead of + # reading them from a precomputed cache. `target_cache_path` is ignored. + online_target=True, + target_cache_path=None, + # JSONL conversation files (same format consumed by prepare_target_cache.py). + train_data_paths=[ + os.path.expanduser("~/.cache/deepspec/train_data/train.jsonl"), + ], + # Drop micro-batches whose only sample has fewer than this many loss tokens. + min_loss_tokens=14, + chat_template="qwen", + max_length=4096, + num_workers=4, +) + + +def finalize_cfg(cfg): + logging_cfg = dict(cfg["logging"]) + project_name = str(cfg["project_name"]) + exp_name = str(cfg["exp_name"]) + logging_cfg["checkpoint_dir"] = os.path.join(BASE_CKPT_DIR, project_name, exp_name) + logging_cfg["tensorboard_dir"] = os.path.join(BASE_TB_DIR, project_name, exp_name) + cfg["logging"] = logging_cfg + + return cfg diff --git a/deepspec/data/__init__.py b/deepspec/data/__init__.py index 722365b3..cc28f48d 100644 --- a/deepspec/data/__init__.py +++ b/deepspec/data/__init__.py @@ -1,3 +1,4 @@ +from .jsonl_dataset import JsonLineDataset from .parser import TEMPLATE_REGISTRY from .target_cache_dataset import ( CacheCollator, @@ -10,6 +11,7 @@ "CacheCollator", "CacheDataset", "ConversationCollator", + "JsonLineDataset", "TEMPLATE_REGISTRY", "validate_train_cache", ] diff --git a/deepspec/data/cuda_prefetcher.py b/deepspec/data/cuda_prefetcher.py index 69eac905..3c57ffac 100644 --- a/deepspec/data/cuda_prefetcher.py +++ b/deepspec/data/cuda_prefetcher.py @@ -34,12 +34,21 @@ def __iter__(self): return self def _fetch_and_transfer(self): - """Pop the next CPU batch from the DataLoader and queue H2D on the side stream.""" - try: - cpu_batch = next(self._iter) - except StopIteration: - self._done = True - return + """Pop the next CPU batch from the DataLoader and queue H2D on the side stream. + + A collator may return ``None`` when every sample in a micro-batch is + filtered out (e.g. ConversationCollator dropping samples below + ``min_loss_tokens``). Skip such batches so downstream consumers never + observe ``None``. + """ + while True: + try: + cpu_batch = next(self._iter) + except StopIteration: + self._done = True + return + if cpu_batch is not None: + break with torch.cuda.stream(self.stream): self._gpu_batch = move_batch_to_device(cpu_batch, self.device) diff --git a/deepspec/modeling/target_extract.py b/deepspec/modeling/target_extract.py new file mode 100644 index 00000000..1e8663e7 --- /dev/null +++ b/deepspec/modeling/target_extract.py @@ -0,0 +1,100 @@ +"""Target-model hidden-state extraction shared by cache generation and online training. + +The cache-generation pipeline (``scripts/data/prepare_target_cache.py``) and the +online training path both need to run the target backbone and capture the hidden +states at a fixed set of layers. That logic lives here so both call sites stay in +lockstep -- an online training step produces exactly the same tensors that would +have been written to (and later read from) the target cache. +""" + +from dataclasses import dataclass + +import torch + + +@dataclass(frozen=True) +class TargetForwardResult: + target_hidden_states: torch.Tensor + target_last_hidden_states: torch.Tensor + + +def get_target_backbone(target_model): + model_type = str(target_model.config.model_type) + if model_type in ("gemma4", "gemma4_unified"): + if hasattr(target_model, "language_model"): + return target_model.language_model + if hasattr(target_model, "model") and hasattr(target_model.model, "language_model"): + return target_model.model.language_model + assert False, "Gemma4 target model must expose a text language_model." + return getattr(target_model, "model", target_model) + + +def get_target_hidden_size(target_model) -> int: + model_type = str(target_model.config.model_type) + if model_type in ("gemma4", "gemma4_unified"): + return int(target_model.config.text_config.hidden_size) + return int(target_model.config.hidden_size) + + +def _get_hook_tensor(output): + if isinstance(output, torch.Tensor): + return output + if isinstance(output, (tuple, list)) and output: + first = output[0] + if isinstance(first, torch.Tensor): + return first + raise TypeError(f"Unsupported target hook output type: {type(output)!r}") + + +def run_target_forward_with_hooks( + *, + target_model, + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + target_layer_ids, +): + backbone = get_target_backbone(target_model) + layer_modules = backbone.layers + target_layer_ids = [int(layer_id) for layer_id in target_layer_ids] + captured_hidden_states = {} + handles = [] + + def capture_layer(layer_id: int): + def hook(_module, _inputs, output): + captured_hidden_states[layer_id] = _get_hook_tensor(output).detach() + + return hook + + try: + if -1 in target_layer_ids: + handles.append( + backbone.embed_tokens.register_forward_hook(capture_layer(-1)) + ) + for layer_id in target_layer_ids: + if layer_id < 0: + continue + handles.append( + layer_modules[layer_id].register_forward_hook(capture_layer(layer_id)) + ) + + with torch.no_grad(): + target_output = target_model( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=False, + use_cache=False, + ) + target_last_hidden_states = target_output.last_hidden_state.detach() + target_hidden_states = torch.cat( + [captured_hidden_states[layer_id] for layer_id in target_layer_ids], + dim=-1, + ) + finally: + for handle in handles: + handle.remove() + captured_hidden_states.clear() + + return TargetForwardResult( + target_hidden_states=target_hidden_states, + target_last_hidden_states=target_last_hidden_states, + ) diff --git a/deepspec/trainer/base_trainer.py b/deepspec/trainer/base_trainer.py index 6b2eef1d..8456216f 100644 --- a/deepspec/trainer/base_trainer.py +++ b/deepspec/trainer/base_trainer.py @@ -10,8 +10,14 @@ from torch.utils.data import DataLoader from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer -from deepspec.data import CacheDataset, validate_train_cache +from deepspec.data import ( + CacheDataset, + ConversationCollator, + JsonLineDataset, + validate_train_cache, +) from deepspec.data.cuda_prefetcher import CUDAPrefetcher +from deepspec.modeling.target_extract import get_target_backbone from deepspec.utils import ( BF16Optimizer, StatelessResumableDistributedSampler, @@ -158,6 +164,9 @@ def __init__(self, local_rank, args): ) self.suspend_controller = SuspendController(device=self.device) self.next_micro_step = 0 + self._online_target = bool(getattr(self.args.data, "online_target", False)) + self.target_backbone = None + self._online_collator = None if is_global_main_process(): ensure_dir(self.checkpoint_dir_root) training_logger.init( @@ -180,12 +189,23 @@ def __init__(self, local_rank, args): self.model = torch.compile(self.model, dynamic=True) self.model = self._wrap_with_fsdp(self.model) - self.train_dataset = CacheDataset(cache_dir=self.args.data.target_cache_path) - validate_train_cache( - train_dataset=self.train_dataset, - draft_model=self.draft_model, - target_model_name_or_path=self.args.model.target_model_name_or_path, - ) + if self._online_target: + self.train_dataset = JsonLineDataset( + data_paths=list(self.args.data.train_data_paths) + ) + self._online_collator = ConversationCollator( + tokenizer=self.tokenizer, + chat_template=self.args.data.chat_template, + max_length=int(self.args.data.max_length), + min_loss_tokens=int(self.args.data.min_loss_tokens), + ) + else: + self.train_dataset = CacheDataset(cache_dir=self.args.data.target_cache_path) + validate_train_cache( + train_dataset=self.train_dataset, + draft_model=self.draft_model, + target_model_name_or_path=self.args.model.target_model_name_or_path, + ) ( self.gradient_accumulation_steps, @@ -257,12 +277,14 @@ def build_models(self): ) draft_model = draft_model.to(device=self.device, dtype=self.precision_dtype) - # Training only uses the target checkpoint to initialize frozen draft - # embeddings and lm_head weights. + # Offline (cached) training only uses the target checkpoint to initialize + # frozen draft embeddings and lm_head weights, then discards it. Online + # training instead keeps the frozen target backbone resident so hidden + # states can be recomputed at every step. target_model = AutoModelForCausalLM.from_pretrained( model_args.target_model_name_or_path, dtype=self.precision_dtype, - ).to(device="cpu").eval() + ).eval() target_embed_tokens = target_model.get_input_embeddings() target_lm_head = target_model.get_output_embeddings() assert (target_lm_head is not None) and (target_embed_tokens is not None) @@ -271,7 +293,18 @@ def build_models(self): lm_head=target_lm_head, freeze=True, ) - del target_model + + if self._online_target: + # Keep the *backbone* (its forward returns last_hidden_state, which + # run_target_forward_with_hooks relies on). Frozen, eval, on-device. + self.target_backbone = get_target_backbone(target_model).to( + device=self.device + ).eval() + for param in self.target_backbone.parameters(): + param.requires_grad_(False) + else: + self.target_backbone = None + del target_model return draft_model, tokenizer def _build_draft_model(self, *, target_config, model_args): @@ -298,7 +331,7 @@ def _build_train_dataloader(self, start_offset_samples=0, num_samples=None): self.train_dataset, batch_size=int(self.args.train.local_batch_size), sampler=sampler, - collate_fn=self.data_collator_cls(), + collate_fn=self._online_collator or self.data_collator_cls(), num_workers=int(self.args.data.num_workers), pin_memory=True, drop_last=True, diff --git a/deepspec/trainer/dspark_trainer.py b/deepspec/trainer/dspark_trainer.py index 487e99a2..0af764ba 100644 --- a/deepspec/trainer/dspark_trainer.py +++ b/deepspec/trainer/dspark_trainer.py @@ -8,6 +8,7 @@ from deepspec.modeling.dspark.qwen3.config import ( build_draft_config as build_qwen3_draft_config, ) +from deepspec.modeling.target_extract import run_target_forward_with_hooks from deepspec.trainer.base_trainer import BaseTrainer @@ -23,6 +24,22 @@ def _build_draft_model(self, *, target_config, model_args): # Training step. def run_batch(self, batch): + if self.target_backbone is not None: + # Online mode: recompute the target hidden states this step instead + # of reading them from the cache. Frozen + no_grad inside the helper. + target_result = run_target_forward_with_hooks( + target_model=self.target_backbone, + input_ids=batch["input_ids"], + attention_mask=batch["attention_mask"], + target_layer_ids=self.draft_model.target_layer_ids, + ) + batch["target_hidden_states"] = target_result.target_hidden_states.to( + self.precision_dtype + ) + batch["target_last_hidden_states"] = ( + target_result.target_last_hidden_states.to(self.precision_dtype) + ) + outputs = self.model( input_ids=batch["input_ids"], target_hidden_states=batch["target_hidden_states"], diff --git a/deepspec/trainer/eagle3_trainer.py b/deepspec/trainer/eagle3_trainer.py index 70343a3b..c5394716 100644 --- a/deepspec/trainer/eagle3_trainer.py +++ b/deepspec/trainer/eagle3_trainer.py @@ -17,6 +17,13 @@ class Qwen3Eagle3Trainer(BaseTrainer): data_collator_cls = CacheCollator def build_models(self): + # Online target mode is currently implemented for DSpark only. Eagle3 + # captures hidden states with a different (layer_id + 1) offset and + # overrides build_models entirely, so guard against silent misuse. + assert not self._online_target, ( + "online_target is not yet supported for Eagle3 trainers; " + "use a precomputed target cache." + ) model_args = self.args.model tokenizer = AutoTokenizer.from_pretrained( diff --git a/scripts/data/prepare_target_cache.py b/scripts/data/prepare_target_cache.py index d072582b..bedeaada 100644 --- a/scripts/data/prepare_target_cache.py +++ b/scripts/data/prepare_target_cache.py @@ -1,5 +1,4 @@ import argparse -from dataclasses import dataclass import json import os @@ -24,6 +23,12 @@ write_target_cache_manifest, ) from deepspec.data.jsonl_dataset import JsonLineDataset +from deepspec.modeling.target_extract import ( + TargetForwardResult, + get_target_backbone as _get_target_backbone, + get_target_hidden_size as _get_target_hidden_size, + run_target_forward_with_hooks, +) from deepspec.utils import ( CustomJSONEncoder, get_git_diff, @@ -46,94 +51,6 @@ torch.set_float32_matmul_precision("high") -@dataclass(frozen=True) -class TargetForwardResult: - target_hidden_states: torch.Tensor - target_last_hidden_states: torch.Tensor - - -def _get_target_backbone(target_model): - model_type = str(target_model.config.model_type) - if model_type in ("gemma4", "gemma4_unified"): - if hasattr(target_model, "language_model"): - return target_model.language_model - if hasattr(target_model, "model") and hasattr(target_model.model, "language_model"): - return target_model.model.language_model - assert False, "Gemma4 target model must expose a text language_model." - return getattr(target_model, "model", target_model) - - -def _get_target_hidden_size(target_model) -> int: - model_type = str(target_model.config.model_type) - if model_type in ("gemma4", "gemma4_unified"): - return int(target_model.config.text_config.hidden_size) - return int(target_model.config.hidden_size) - - -def _get_hook_tensor(output): - if isinstance(output, torch.Tensor): - return output - if isinstance(output, (tuple, list)) and output: - first = output[0] - if isinstance(first, torch.Tensor): - return first - raise TypeError(f"Unsupported target hook output type: {type(output)!r}") - - -def run_target_forward_with_hooks( - *, - target_model, - input_ids: torch.Tensor, - attention_mask: torch.Tensor, - target_layer_ids, -): - backbone = _get_target_backbone(target_model) - layer_modules = backbone.layers - target_layer_ids = [int(layer_id) for layer_id in target_layer_ids] - captured_hidden_states = {} - handles = [] - - def capture_layer(layer_id: int): - def hook(_module, _inputs, output): - captured_hidden_states[layer_id] = _get_hook_tensor(output).detach() - - return hook - - try: - if -1 in target_layer_ids: - handles.append( - backbone.embed_tokens.register_forward_hook(capture_layer(-1)) - ) - for layer_id in target_layer_ids: - if layer_id < 0: - continue - handles.append( - layer_modules[layer_id].register_forward_hook(capture_layer(layer_id)) - ) - - with torch.no_grad(): - target_output = target_model( - input_ids=input_ids, - attention_mask=attention_mask, - output_hidden_states=False, - use_cache=False, - ) - target_last_hidden_states = target_output.last_hidden_state.detach() - target_hidden_states = torch.cat( - [captured_hidden_states[layer_id] for layer_id in target_layer_ids], - dim=-1, - ) - finally: - for handle in handles: - handle.remove() - captured_hidden_states.clear() - - return TargetForwardResult( - target_hidden_states=target_hidden_states, - target_last_hidden_states=target_last_hidden_states, - ) - - def parse_args(): parser = argparse.ArgumentParser() parser.add_argument("--config", required=True) diff --git a/tests/test_online_target_extract.py b/tests/test_online_target_extract.py new file mode 100644 index 00000000..077c5bf8 --- /dev/null +++ b/tests/test_online_target_extract.py @@ -0,0 +1,165 @@ +"""CPU correctness tests for online target hidden-state extraction. + +These validate the core claim behind online training mode: that recomputing +target hidden states at train time with ``run_target_forward_with_hooks`` +produces exactly the tensors the offline cache pipeline would have stored. + +Runs on CPU with a tiny randomly-initialized Qwen3 model -- no GPU, no +downloads. Run with: ``pytest tests/test_online_target_extract.py`` +""" + +import importlib.util +import os + +import torch +from transformers import Qwen3Config, Qwen3ForCausalLM + +# Load deepspec/modeling/target_extract.py directly. Importing it via the +# package would execute deepspec/modeling/__init__.py, which eagerly imports +# the Gemma4 modeling classes -- those require a newer transformers than some +# environments have installed. The module under test has no such dependency. +_MODULE_PATH = os.path.join( + os.path.dirname(os.path.dirname(__file__)), + "deepspec", + "modeling", + "target_extract.py", +) +_spec = importlib.util.spec_from_file_location("target_extract", _MODULE_PATH) +_target_extract = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(_target_extract) + +get_target_backbone = _target_extract.get_target_backbone +run_target_forward_with_hooks = _target_extract.run_target_forward_with_hooks + + +def _tiny_model(): + torch.manual_seed(0) + config = Qwen3Config( + vocab_size=256, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=6, + num_attention_heads=4, + num_key_value_heads=2, + max_position_embeddings=64, + ) + # The helper consumes a *backbone* (returns last_hidden_state), exactly as + # cache generation loads it via AutoModel and as build_models extracts it + # via get_target_backbone. Use the inner .model here. + return Qwen3ForCausalLM(config).eval().model + + +def _reference_capture(model, input_ids, attention_mask, target_layer_ids): + """Independent re-implementation of the capture, mirroring cache-gen.""" + backbone = get_target_backbone(model) + captured = {} + handles = [] + + def hook(layer_id): + def _hook(_m, _i, output): + t = output[0] if isinstance(output, (tuple, list)) else output + captured[layer_id] = t.detach() + return _hook + + for lid in target_layer_ids: + handles.append(backbone.layers[lid].register_forward_hook(hook(lid))) + try: + with torch.no_grad(): + out = model( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=False, + use_cache=False, + ) + last = out.last_hidden_state.detach() + hidden = torch.cat([captured[lid] for lid in target_layer_ids], dim=-1) + finally: + for h in handles: + h.remove() + return hidden, last + + +def test_online_extraction_matches_reference(): + model = _tiny_model() + target_layer_ids = [1, 3, 5] + input_ids = torch.randint(0, 256, (2, 16)) + attention_mask = torch.ones_like(input_ids) + + result = run_target_forward_with_hooks( + target_model=model, + input_ids=input_ids, + attention_mask=attention_mask, + target_layer_ids=target_layer_ids, + ) + ref_hidden, ref_last = _reference_capture( + model, input_ids, attention_mask, target_layer_ids + ) + + assert result.target_hidden_states.shape == (2, 16, 32 * len(target_layer_ids)) + assert result.target_last_hidden_states.shape == (2, 16, 32) + torch.testing.assert_close(result.target_hidden_states, ref_hidden) + torch.testing.assert_close(result.target_last_hidden_states, ref_last) + + +def test_padding_does_not_change_real_token_hidden_states(): + """A right-padded batch must yield the same hidden states on the real + tokens as an unpadded single-sequence forward. This is the padding/masking + risk flagged during design review.""" + model = _tiny_model() + target_layer_ids = [1, 3, 5] + + real_len = 10 + seq = torch.randint(1, 256, (1, real_len)) + + # Unpadded reference. + unpadded = run_target_forward_with_hooks( + target_model=model, + input_ids=seq, + attention_mask=torch.ones_like(seq), + target_layer_ids=target_layer_ids, + ) + + # Right-padded to length 16 with a masked-out pad region. + pad_len = 16 + padded_ids = torch.zeros((1, pad_len), dtype=torch.long) + padded_ids[0, :real_len] = seq[0] + mask = torch.zeros((1, pad_len), dtype=torch.long) + mask[0, :real_len] = 1 + padded = run_target_forward_with_hooks( + target_model=model, + input_ids=padded_ids, + attention_mask=mask, + target_layer_ids=target_layer_ids, + ) + + torch.testing.assert_close( + padded.target_hidden_states[:, :real_len], + unpadded.target_hidden_states, + rtol=1e-4, + atol=1e-4, + ) + torch.testing.assert_close( + padded.target_last_hidden_states[:, :real_len], + unpadded.target_last_hidden_states, + rtol=1e-4, + atol=1e-4, + ) + + +def test_target_backbone_unwraps_causal_lm(): + # A ForCausalLM must unwrap to its .model backbone (the layer container). + full = Qwen3ForCausalLM( + Qwen3Config( + vocab_size=256, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + ) + ).eval() + backbone = get_target_backbone(full) + assert backbone is full.model + assert hasattr(backbone, "layers") + # And an already-unwrapped backbone is returned as-is. + assert get_target_backbone(full.model) is full.model From f4cfbde5b99bbd8834e6bcbe38605f71c188c4b7 Mon Sep 17 00:00:00 2001 From: Ofir Ben Shoham Date: Tue, 30 Jun 2026 11:03:36 +0000 Subject: [PATCH 2/2] Document online training; keep target model on CPU in offline mode - README: add an Online Training section covering the no-cache path, the compute/memory tradeoff, the data.online_target config, and the JSONL input format. Note Eagle3 online training is not yet supported. - train.sh: accept an optional config-path argument (defaults to the cached DSpark Qwen3-4B config) so online configs can be launched without editing the script; target_cache_path is ignored online. - base_trainer: restore .to(device="cpu") on the target model load so offline (cached) training never places the soon-discarded target model on the GPU. Online still moves only the backbone to the device. --- README.md | 50 ++++++++++++++++++++++++++++++++ deepspec/trainer/base_trainer.py | 5 +++- scripts/train/train.sh | 9 +++++- 3 files changed, 62 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index dbb79990..0d871755 100644 --- a/README.md +++ b/README.md @@ -20,6 +20,10 @@ Run the stages in order — each stage's output feeds the next: 2. **Training** — train a draft model against the cached target outputs. 3. **Evaluation** — measure speculative-decoding acceptance on benchmark tasks. +> **Tip:** [Online training](#online-training-no-target-cache) can skip the target-cache +> step entirely (Data Preparation steps 1–2 only) by recomputing target hidden states +> during training. This trades extra GPU compute per step for not storing the cache. + ## Data Preparation See [scripts/data/README.md](./scripts/data/README.md) for the step-by-step data pipeline: @@ -38,6 +42,52 @@ bash scripts/train/train.sh Hardware: the default configs and scripts assume a single node with 8 GPUs. For fewer GPUs, reduce `CUDA_VISIBLE_DEVICES`. +### Online Training (no target cache) + +The default training path reads target hidden states from a precomputed target cache, +which can be very large (roughly 38 TB for `Qwen/Qwen3-4B`; see the +[storage warning](./scripts/data/README.md#step-3-prepare-target-cache)). **Online +training** avoids that cache: it keeps the frozen target model resident on the GPU and +recomputes the target hidden states for each batch during training. This trades extra +GPU compute and memory per step for not having to generate or store the cache, which is +convenient for smaller datasets or storage-constrained machines. + +To use it, point `train.sh` at an online config — currently +[config/dspark/dspark_qwen3_4b_online.py](./config/dspark/dspark_qwen3_4b_online.py): + +```bash +bash scripts/train/train.sh \ + config/dspark/dspark_qwen3_4b_online.py +``` + +(or run `train.py --config config/dspark/dspark_qwen3_4b_online.py` directly). + +An online config differs from its cached counterpart only in the `data` block: + +```python +data = dict( + # Recompute target hidden states each step instead of reading a cache. + online_target=True, + target_cache_path=None, # ignored in online mode + # JSONL conversation files (same format as the cache pipeline's input; + # see scripts/data/README.md). Each line: {"conversations": [{"role", "content"}, ...]}. + train_data_paths=["~/.cache/deepspec/train_data/train.jsonl"], + min_loss_tokens=14, # drop micro-batches with too few loss tokens + chat_template="qwen", + max_length=4096, + num_workers=4, +) +``` + +The input JSONL is the same format produced by Data Preparation **steps 1–2** (the +regenerated answers from `generate_train_data.py`); you skip **step 3** (cache +generation). Online mode is only implemented for the DSpark trainer — Eagle3 online +training is not yet supported. + +> **Note:** because the full target model stays on the GPU and runs a forward pass every +> step, online training uses more GPU memory and is slower per step than reading from the +> cache. If you hit OOM, lower `train.local_batch_size` or `data.max_length`. + ## Evaluation diff --git a/deepspec/trainer/base_trainer.py b/deepspec/trainer/base_trainer.py index 8456216f..c57f2253 100644 --- a/deepspec/trainer/base_trainer.py +++ b/deepspec/trainer/base_trainer.py @@ -281,10 +281,13 @@ def build_models(self): # frozen draft embeddings and lm_head weights, then discards it. Online # training instead keeps the frozen target backbone resident so hidden # states can be recomputed at every step. + # Load on CPU first. Offline (cached) training discards this model after + # copying embeddings, so it must never occupy GPU memory. Online training + # moves only the backbone to the GPU below. target_model = AutoModelForCausalLM.from_pretrained( model_args.target_model_name_or_path, dtype=self.precision_dtype, - ).eval() + ).to(device="cpu").eval() target_embed_tokens = target_model.get_input_embeddings() target_lm_head = target_model.get_output_embeddings() assert (target_lm_head is not None) and (target_embed_tokens is not None) diff --git a/scripts/train/train.sh b/scripts/train/train.sh index 05a06188..600e32ae 100644 --- a/scripts/train/train.sh +++ b/scripts/train/train.sh @@ -27,6 +27,11 @@ export WORLD_SIZE=${WORLD_SIZE:-1} # config/eagle3/eagle3_qwen3_8b.py # config/eagle3/eagle3_qwen3_14b.py +# The training config. Defaults to the cached DSpark Qwen3-4B config; override it by +# passing a config path as the first argument, e.g. to use online training: +# bash scripts/train/train.sh config/dspark/dspark_qwen3_4b_online.py +config_path=${1:-config/dspark/dspark_qwen3_4b.py} + target_cache_dir=${target_cache_dir:-${HOME}/.cache/deepspec/qwen3_4b_target_cache} # --opts overrides any config field by dotted key path: --opts "=". @@ -40,6 +45,8 @@ target_cache_dir=${target_cache_dir:-${HOME}/.cache/deepspec/qwen3_4b_target_cac # with more memory (e.g. 4 or 8 on 80GB cards), or keep it at 1 if you hit OOM. # Override it without editing the config via: # --opts "train.local_batch_size=4" +# data.target_cache_path is ignored by online configs (data.online_target=True), which +# stream conversations from data.train_data_paths instead of reading a target cache. python train.py \ - --config config/dspark/dspark_qwen3_4b.py \ + --config "${config_path}" \ --opts "data.target_cache_path=${target_cache_dir}"