Add online training mode for DSpark (no precomputed target cache)

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
 

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

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",
 ]

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)
 

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,
+    )

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,8 +277,13 @@ 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.
+        # 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,
@@ -271,7 +296,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 +334,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,