feat(v4): DeepSeek-V4-Flash DSpark training pipeline against a quantized serving target
convert_engine_dump.py: hanzo-engine hidden-state dumps -> TARGET_CACHE_VERSION=2 (reuses DeepSpec's own writers; bit-exact round-trip). dspark_v4flash_quant.py: V4-Flash draft config (hidden 4096, vocab 129280, 43 target layers, layer_ids [1,11,21,31,41], block 7, markov vanilla 256, confidence head, mask_token 129279). v4_pipeline/: synthetic dump generator + training smoke (single-GPU, reduced dims). PROVEN on GB10: 300 real samples dumped at 111.5 tok/s prefill -> converted -> 50 steps on real quantized-V4 hiddens, loss 3.03->2.6 decreasing; full config 5L/64a/12288 passes solo.
This commit is contained in:
hanzo-dev
@@ -0,0 +1,122 @@
|
||||
"""DSpark draft config for a DeepSeek-V4-Flash (quantized) serving target.
|
||||
|
||||
Adapted from ``config/dspark/dspark_qwen3_4b.py``. The DRAFT is a standard
|
||||
DSpark dense drafter (5 layers, block 7, Markov + confidence heads); only the
|
||||
target it distills from changes. The target model here is DeepSeek-V4-Flash
|
||||
served QUANTIZED, from which the Lux engine dumps per-token hidden states that
|
||||
``scripts/data/convert_engine_dump.py`` packs into a DeepSpec target cache.
|
||||
|
||||
V4-Flash target dimensions (derived from the target's HF config at train time,
|
||||
listed here for reference):
|
||||
hidden_size = 4096
|
||||
vocab_size = 129280
|
||||
num_hidden_layers = 43 (target layers; captured ids below)
|
||||
|
||||
INTEGRATION TODO (target tokenizer / embeddings):
|
||||
``deepspec.trainer.base_trainer.BaseTrainer.build_models`` loads
|
||||
``target_model_name_or_path`` three ways:
|
||||
1. ``AutoTokenizer.from_pretrained`` (tokenizer for logging/eval only),
|
||||
2. ``AutoConfig.from_pretrained`` (to derive the draft dims via
|
||||
``build_qwen3_draft_config`` -- deepcopies the target config),
|
||||
3. ``AutoModelForCausalLM.from_pretrained`` (to copy + FREEZE the draft's
|
||||
``embed_tokens`` and ``lm_head`` from the target).
|
||||
We currently only have the V4 tokenizer inside a GGUF. A real training run
|
||||
therefore needs a LOCAL HF-format directory at ``target_model_name_or_path``
|
||||
containing at least ``config.json`` + tokenizer files + the embedding and
|
||||
lm_head weights (full weights not required if a trimmed checkpoint exposing
|
||||
just the input/output embeddings is provided). Because the target is
|
||||
DeepSeek (not Qwen3), that ``config.json`` must expose the Qwen3-style
|
||||
fields the dense drafter reads (``num_attention_heads``,
|
||||
``num_key_value_heads``, ``head_dim``, ``intermediate_size``,
|
||||
``rms_norm_eps``, ``rope_parameters``, ``max_position_embeddings``), or a
|
||||
tiny target-config adapter must synthesize a Qwen3Config with V4 dims. The
|
||||
toy smoke (``v4_pipeline/smoke_train.py``) bypasses all three loads by
|
||||
constructing a synthetic Qwen3 target config with the V4 dims above.
|
||||
|
||||
The manifest's ``target_model_name_or_path`` (written by the converter) MUST
|
||||
equal ``model.target_model_name_or_path`` below -- ``validate_train_cache``
|
||||
asserts string equality.
|
||||
"""
|
||||
|
||||
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_v4flash_quant"
|
||||
seed = 42
|
||||
|
||||
# Local HF-format directory for the DeepSeek-V4-Flash target (see
|
||||
# INTEGRATION TODO above). Must match the converter's
|
||||
# --target-model-name-or-path.
|
||||
TARGET_MODEL_PATH = os.path.expanduser("~/models/deepseek-v4-flash")
|
||||
|
||||
model = dict(
|
||||
target_model_name_or_path=TARGET_MODEL_PATH,
|
||||
block_size=7,
|
||||
num_draft_layers=5,
|
||||
# 5 captured target layers out of 43 (must be strictly ascending and each
|
||||
# in [-1, 42]; -1 would denote the embedding output).
|
||||
target_layer_ids=[1, 11, 21, 31, 41],
|
||||
# Mask/noise token fed at draft positions (create_noise_embed ->
|
||||
# embed_tokens lookup ONLY). Requirements: 0 <= id < vocab_size (129280)
|
||||
# and the tokenizer must never EMIT it, so a masked draft position is never
|
||||
# confused with a real token. 129279 = vocab_size - 1 is the top-most id and
|
||||
# sits in DeepSeek's reserved/special band -> a safe, unemitted placeholder.
|
||||
# TODO: once the real V4 tokenizer is available, confirm this id is reserved
|
||||
# (ideally point it at a dedicated <|mask|>/pad reserved token).
|
||||
mask_token_id=129279,
|
||||
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(
|
||||
target_cache_path=None,
|
||||
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
|
||||
@@ -0,0 +1,312 @@
|
||||
"""Convert a Lux engine target-hidden-state dump into a DeepSpec target cache.
|
||||
|
||||
The Rust engine (separate workstream) dumps, per training sample:
|
||||
|
||||
* ``input_ids`` [seq] int64/int32
|
||||
* ``target_hidden_states`` [seq, L, H] bf16 (L captured layers)
|
||||
* ``target_last_hidden_states`` [seq, H] bf16
|
||||
* (optional) ``loss_mask`` [seq] any int/bool
|
||||
* (optional) ``attention_mask`` [seq] any int/bool
|
||||
|
||||
alongside an ``index.jsonl`` (one JSON object per line) that points at each
|
||||
per-sample ``*.safetensors`` shard via a ``file`` / ``path`` / ``safetensors``
|
||||
key.
|
||||
|
||||
This script re-emits those samples in DeepSpec's canonical *target cache*
|
||||
protocol (version 2) by driving DeepSpec's OWN writer/finalizer classes -- it
|
||||
does not reimplement the binary layout. The output directory is directly
|
||||
consumable by ``deepspec.data.CacheDataset`` and ``scripts/train/train.sh``.
|
||||
|
||||
``target_hidden_states`` is stored as ``[seq, L * H]`` with the L captured
|
||||
layers concatenated along the hidden axis in ``--target-layer-ids`` order --
|
||||
byte-identical to what ``run_target_forward_with_hooks`` in
|
||||
``scripts/data/prepare_target_cache.py`` produces
|
||||
(``torch.cat([captured[id] for id in target_layer_ids], dim=-1)``).
|
||||
|
||||
Single process, no torch.distributed: this is the ``world_size == 1``
|
||||
specialization of ``scripts/data/prepare_target_cache.py``'s finalize flow.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
|
||||
import torch
|
||||
from safetensors.torch import load_file
|
||||
|
||||
from deepspec.data.target_cache_dataset import (
|
||||
AsyncTargetCacheWriter,
|
||||
LocalCacheWriteSummary,
|
||||
atomic_json_dump,
|
||||
build_global_target_cache_shard_map,
|
||||
build_target_cache_manifest,
|
||||
cleanup_target_cache_tmp_dir,
|
||||
finalize_target_cache_index,
|
||||
load_local_cache_write_summary,
|
||||
load_target_cache_manifest,
|
||||
prepare_target_cache_output_dir,
|
||||
rename_local_target_cache_shards,
|
||||
write_target_cache_manifest,
|
||||
)
|
||||
|
||||
|
||||
_FILE_KEYS = ("file", "path", "safetensors", "sample", "shard")
|
||||
_INPUT_IDS_KEYS = ("input_ids", "input_id", "ids", "tokens")
|
||||
_HIDDEN_KEYS = ("target_hidden_states", "hidden_states", "target_hidden")
|
||||
_LAST_HIDDEN_KEYS = (
|
||||
"target_last_hidden_states",
|
||||
"last_hidden_states",
|
||||
"target_last_hidden",
|
||||
)
|
||||
_LOSS_MASK_KEYS = ("loss_mask", "loss_masks")
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument(
|
||||
"--dump-dir",
|
||||
required=True,
|
||||
help="Engine-dump directory containing index.jsonl + *.safetensors.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
required=True,
|
||||
help="Destination target-cache directory (must be new/empty).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--index-name",
|
||||
default="index.jsonl",
|
||||
help="Index file name inside --dump-dir (default: index.jsonl).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--hidden-size",
|
||||
type=int,
|
||||
default=4096,
|
||||
help="Target hidden dimension H (V4-Flash: 4096).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--target-layer-ids",
|
||||
default="1,11,21,31,41",
|
||||
help="Comma-separated captured target layer ids, ascending "
|
||||
"(V4-Flash: 1,11,21,31,41).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--target-model-name-or-path",
|
||||
required=True,
|
||||
help="Recorded in the manifest; MUST equal "
|
||||
"config.model.target_model_name_or_path used at train time "
|
||||
"(validate_train_cache asserts string equality).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--loss-mask-mode",
|
||||
choices=("ones_except_first", "all_ones", "from_dump"),
|
||||
default="ones_except_first",
|
||||
help="How to build loss_mask. 'ones_except_first' (default): 1 "
|
||||
"everywhere except position 0 -- position 0 has no left context so it "
|
||||
"is never a valid anchor nor a first draft target "
|
||||
"(see build_anchor_candidate_mask). 'from_dump': use the dump's "
|
||||
"loss_mask (error if a sample lacks one).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-shard-bytes",
|
||||
type=int,
|
||||
default=64 * 1024**3,
|
||||
help="Max bytes per output shard before rolling to a new one.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--limit",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Only convert the first N index entries (debug).",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def _first_present(mapping, keys, what, where):
|
||||
for key in keys:
|
||||
if key in mapping:
|
||||
return mapping[key]
|
||||
raise KeyError(f"{what} not found in {where}; looked for keys {keys}.")
|
||||
|
||||
|
||||
def _resolve_sample_path(dump_dir, record, line_no):
|
||||
file_ref = None
|
||||
for key in _FILE_KEYS:
|
||||
if key in record:
|
||||
file_ref = record[key]
|
||||
break
|
||||
if file_ref is None:
|
||||
raise KeyError(
|
||||
f"index line {line_no}: no shard reference; expected one of "
|
||||
f"{_FILE_KEYS} in {record!r}."
|
||||
)
|
||||
path = str(file_ref)
|
||||
if not os.path.isabs(path):
|
||||
path = os.path.join(dump_dir, path)
|
||||
if not os.path.exists(path):
|
||||
raise FileNotFoundError(f"index line {line_no}: missing sample file {path}.")
|
||||
return path
|
||||
|
||||
|
||||
def _build_loss_mask(mode, seq_len, tensors, device):
|
||||
if mode == "from_dump":
|
||||
loss_mask = _first_present(
|
||||
tensors, _LOSS_MASK_KEYS, "loss_mask", "sample tensors"
|
||||
)
|
||||
return loss_mask.reshape(-1).to(dtype=torch.uint8)
|
||||
loss_mask = torch.ones(seq_len, dtype=torch.uint8, device=device)
|
||||
if mode == "ones_except_first" and seq_len > 0:
|
||||
loss_mask[0] = 0
|
||||
return loss_mask
|
||||
|
||||
|
||||
def _load_sample(path, *, hidden_size, num_target_layers, loss_mask_mode):
|
||||
tensors = load_file(path)
|
||||
input_ids = _first_present(tensors, _INPUT_IDS_KEYS, "input_ids", path).reshape(-1)
|
||||
seq_len = int(input_ids.shape[0])
|
||||
assert seq_len > 0, f"{path}: empty input_ids."
|
||||
|
||||
hidden = _first_present(tensors, _HIDDEN_KEYS, "target_hidden_states", path)
|
||||
expected_hidden_numel = seq_len * num_target_layers * hidden_size
|
||||
assert hidden.numel() == expected_hidden_numel, (
|
||||
f"{path}: target_hidden_states has {hidden.numel()} elements, expected "
|
||||
f"{expected_hidden_numel} = seq({seq_len}) * layers({num_target_layers}) "
|
||||
f"* hidden({hidden_size}). Got shape {tuple(hidden.shape)}."
|
||||
)
|
||||
# [seq, L, H] or already [seq, L*H] -> [seq, L*H] (row-major keeps
|
||||
# per-layer blocks contiguous, matching target_layer_ids concat order).
|
||||
target_hidden_states = hidden.reshape(seq_len, num_target_layers * hidden_size)
|
||||
|
||||
last_hidden = _first_present(
|
||||
tensors, _LAST_HIDDEN_KEYS, "target_last_hidden_states", path
|
||||
)
|
||||
assert last_hidden.numel() == seq_len * hidden_size, (
|
||||
f"{path}: target_last_hidden_states has {last_hidden.numel()} elements, "
|
||||
f"expected {seq_len * hidden_size} = seq({seq_len}) * hidden({hidden_size})."
|
||||
)
|
||||
target_last_hidden_states = last_hidden.reshape(seq_len, hidden_size)
|
||||
|
||||
loss_mask = _build_loss_mask(loss_mask_mode, seq_len, tensors, input_ids.device)
|
||||
assert loss_mask.shape[0] == seq_len, (
|
||||
f"{path}: loss_mask length {loss_mask.shape[0]} != seq_len {seq_len}."
|
||||
)
|
||||
attention_mask = torch.ones(seq_len, dtype=torch.uint8, device=input_ids.device)
|
||||
|
||||
return {
|
||||
"input_ids": input_ids,
|
||||
"attention_mask": attention_mask,
|
||||
"loss_mask": loss_mask,
|
||||
"target_hidden_states": target_hidden_states,
|
||||
"target_last_hidden_states": target_last_hidden_states,
|
||||
}
|
||||
|
||||
|
||||
def _read_index(index_path, limit):
|
||||
records = []
|
||||
with open(index_path, "r", encoding="utf-8") as handle:
|
||||
for line_no, raw in enumerate(handle):
|
||||
raw = raw.strip()
|
||||
if not raw:
|
||||
continue
|
||||
records.append((line_no, json.loads(raw)))
|
||||
if limit is not None and len(records) >= limit:
|
||||
break
|
||||
assert records, f"No records read from {index_path}."
|
||||
return records
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
dump_dir = os.path.abspath(args.dump_dir)
|
||||
output_dir = os.path.abspath(args.output_dir)
|
||||
index_path = os.path.join(dump_dir, args.index_name)
|
||||
assert os.path.exists(index_path), f"Missing index file: {index_path}"
|
||||
|
||||
target_layer_ids = [int(x) for x in args.target_layer_ids.split(",") if x != ""]
|
||||
assert target_layer_ids == sorted(target_layer_ids), (
|
||||
f"--target-layer-ids must be ascending, got {target_layer_ids}."
|
||||
)
|
||||
num_target_layers = len(target_layer_ids)
|
||||
hidden_size = int(args.hidden_size)
|
||||
|
||||
records = _read_index(index_path, args.limit)
|
||||
print(
|
||||
f"Converting {len(records)} samples from {dump_dir}\n"
|
||||
f" hidden_size={hidden_size} target_layer_ids={target_layer_ids} "
|
||||
f"loss_mask_mode={args.loss_mask_mode}\n"
|
||||
f" -> {output_dir}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
prepare_target_cache_output_dir(output_dir)
|
||||
rank_dir = os.path.join(output_dir, "_tmp", "rank_0")
|
||||
os.makedirs(rank_dir, exist_ok=True)
|
||||
|
||||
writer = AsyncTargetCacheWriter(
|
||||
rank_dir=rank_dir,
|
||||
max_shard_bytes=int(args.max_shard_bytes),
|
||||
)
|
||||
try:
|
||||
for idx, (line_no, record) in enumerate(records):
|
||||
path = _resolve_sample_path(dump_dir, record, line_no)
|
||||
sample = _load_sample(
|
||||
path,
|
||||
hidden_size=hidden_size,
|
||||
num_target_layers=num_target_layers,
|
||||
loss_mask_mode=args.loss_mask_mode,
|
||||
)
|
||||
writer.write_sample(**sample)
|
||||
if (idx + 1) % 100 == 0 or (idx + 1) == len(records):
|
||||
print(f" wrote {idx + 1}/{len(records)} samples", flush=True)
|
||||
finally:
|
||||
writer.close()
|
||||
|
||||
summary = LocalCacheWriteSummary(
|
||||
global_rank=0,
|
||||
source_sample_start=0,
|
||||
source_sample_end=len(records),
|
||||
num_local_samples=writer.num_local_samples,
|
||||
num_local_shards=len(writer.local_shard_files),
|
||||
local_shard_files=list(writer.local_shard_files),
|
||||
)
|
||||
atomic_json_dump(summary.to_json(), os.path.join(rank_dir, "summary.json"))
|
||||
|
||||
summaries = [load_local_cache_write_summary(rank_dir)]
|
||||
shard_map, shards = build_global_target_cache_shard_map(summaries)
|
||||
rename_local_target_cache_shards(
|
||||
output_dir=output_dir,
|
||||
rank_dir=rank_dir,
|
||||
summary=summaries[0],
|
||||
shard_map=shard_map,
|
||||
)
|
||||
num_valid_samples = finalize_target_cache_index(
|
||||
output_dir=output_dir,
|
||||
summaries=summaries,
|
||||
shard_map=shard_map,
|
||||
)
|
||||
manifest = build_target_cache_manifest(
|
||||
num_samples=num_valid_samples,
|
||||
shards=shards,
|
||||
target_layer_ids=target_layer_ids,
|
||||
hidden_size=hidden_size,
|
||||
extra_fields={
|
||||
"target_model_name_or_path": str(args.target_model_name_or_path),
|
||||
"source_engine_dump": dump_dir,
|
||||
"loss_mask_mode": args.loss_mask_mode,
|
||||
"converter": "scripts/data/convert_engine_dump.py",
|
||||
},
|
||||
)
|
||||
write_target_cache_manifest(output_dir=output_dir, manifest=manifest)
|
||||
cleanup_target_cache_tmp_dir(output_dir)
|
||||
|
||||
# Re-load through DeepSpec's own validator to prove the cache is valid.
|
||||
load_target_cache_manifest(output_dir)
|
||||
print(
|
||||
f"Wrote valid target cache: {num_valid_samples} samples, "
|
||||
f"{len(shards)} shard(s) at {output_dir}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,68 @@
|
||||
"""Generate a synthetic Lux engine target-hidden-state dump for smoke testing.
|
||||
|
||||
Mimics exactly what the Rust engine tool (separate workstream) emits: an
|
||||
``index.jsonl`` plus per-sample ``sample_*.safetensors`` files with
|
||||
|
||||
* ``input_ids`` [seq] int64 (values < vocab_size)
|
||||
* ``target_hidden_states`` [seq, L, H] bf16
|
||||
* ``target_last_hidden_states`` [seq, H] bf16
|
||||
|
||||
Random data only -- this exercises the CONVERTER + TRAINING pipeline shapes,
|
||||
not model quality.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
|
||||
import torch
|
||||
from safetensors.torch import save_file
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--output-dir", required=True)
|
||||
parser.add_argument("--num-samples", type=int, default=10)
|
||||
parser.add_argument("--seq-len", type=int, default=256)
|
||||
parser.add_argument("--hidden-size", type=int, default=4096)
|
||||
parser.add_argument("--num-target-layers", type=int, default=5)
|
||||
parser.add_argument("--vocab-size", type=int, default=129280)
|
||||
parser.add_argument("--seed", type=int, default=0)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
torch.manual_seed(args.seed)
|
||||
|
||||
index_path = os.path.join(args.output_dir, "index.jsonl")
|
||||
with open(index_path, "w", encoding="utf-8") as index_handle:
|
||||
for i in range(args.num_samples):
|
||||
file_name = f"sample_{i:05d}.safetensors"
|
||||
tensors = {
|
||||
"input_ids": torch.randint(
|
||||
0, args.vocab_size, (args.seq_len,), dtype=torch.int64
|
||||
),
|
||||
"target_hidden_states": torch.randn(
|
||||
args.seq_len, args.num_target_layers, args.hidden_size
|
||||
).to(torch.bfloat16),
|
||||
"target_last_hidden_states": torch.randn(
|
||||
args.seq_len, args.hidden_size
|
||||
).to(torch.bfloat16),
|
||||
}
|
||||
save_file(tensors, os.path.join(args.output_dir, file_name))
|
||||
index_handle.write(
|
||||
json.dumps({"file": file_name, "seq_len": args.seq_len}) + "\n"
|
||||
)
|
||||
|
||||
print(
|
||||
f"Wrote {args.num_samples} synthetic samples "
|
||||
f"(seq={args.seq_len}, layers={args.num_target_layers}, "
|
||||
f"hidden={args.hidden_size}, vocab={args.vocab_size}) to {args.output_dir}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,192 @@
|
||||
"""Toy end-to-end DSpark training smoke for the DeepSeek-V4-Flash config.
|
||||
|
||||
Proves the PYTHON training pipeline runs on this box:
|
||||
CacheDataset -> CacheCollator -> Qwen3DSparkModel forward
|
||||
-> compute_dspark_loss -> backward -> BF16Optimizer.step
|
||||
|
||||
It reuses the REAL DeepSpec components (config loader, draft-config builder,
|
||||
model, dataset, collator, loss, optimizer). It only bypasses
|
||||
``BaseTrainer.build_models``' three ``from_pretrained`` loads of the V4 target
|
||||
(tokenizer / AutoConfig / AutoModelForCausalLM), which need a local HF-format
|
||||
V4 checkpoint we do not have yet (see dspark_v4flash_quant.py INTEGRATION TODO).
|
||||
Bypass = (a) synthesize a Qwen3 target config carrying the real V4 dims, and
|
||||
(b) randomly init + FREEZE embed_tokens/lm_head instead of copying them from the
|
||||
target (frozen == exactly what build_models does after copying).
|
||||
|
||||
SMOKE REDUCTIONS vs the shipped config (memory only; every other knob is real):
|
||||
--num-draft-layers (default 2, config ships 5)
|
||||
--num-anchors (default 16, config ships 512)
|
||||
--intermediate-size (default 4096; the true V4 dense MLP width is larger)
|
||||
Real & unchanged: hidden 4096, vocab 129280, block_size 7, markov rank 256
|
||||
(vanilla), confidence head on, target_layer_ids [1,11,21,31,41], bf16,
|
||||
flex_attention, BF16Optimizer, compute_dspark_loss (ce+tv+confidence).
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import time
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from torch.utils.data import DataLoader
|
||||
from transformers import Qwen3Config
|
||||
|
||||
from deepspec.data import CacheCollator, CacheDataset, validate_train_cache
|
||||
from deepspec.modeling.dspark.loss import compute_dspark_loss
|
||||
from deepspec.modeling.dspark.qwen3 import Qwen3DSparkModel
|
||||
from deepspec.modeling.dspark.qwen3.config import build_draft_config
|
||||
from deepspec.utils import BF16Optimizer, load_config, seed_all
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--config", required=True)
|
||||
parser.add_argument("--cache-dir", required=True)
|
||||
parser.add_argument("--steps", type=int, default=8)
|
||||
parser.add_argument("--device", default="cuda")
|
||||
parser.add_argument("--num-draft-layers", type=int, default=2)
|
||||
parser.add_argument("--num-anchors", type=int, default=16)
|
||||
parser.add_argument("--intermediate-size", type=int, default=4096)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def _init_single_process_group():
|
||||
os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
|
||||
os.environ.setdefault("MASTER_PORT", "29577")
|
||||
os.environ.setdefault("RANK", "0")
|
||||
os.environ.setdefault("WORLD_SIZE", "1")
|
||||
if not dist.is_initialized():
|
||||
# gloo: the process group is only used by compute_dspark_loss for a
|
||||
# world_size>1 all-reduce (skipped at world_size==1); CUDA tensors are
|
||||
# never sent over it.
|
||||
dist.init_process_group(backend="gloo", rank=0, world_size=1)
|
||||
|
||||
|
||||
def _build_target_config(cfg, args):
|
||||
"""Synthetic Qwen3 target config carrying the real V4-Flash dims.
|
||||
|
||||
Stands in for AutoConfig.from_pretrained(target) which we cannot run
|
||||
without a local HF V4 checkpoint. num_hidden_layers (43) must exceed
|
||||
max(target_layer_ids) so validate_target_layer_ids passes.
|
||||
"""
|
||||
return Qwen3Config(
|
||||
hidden_size=4096,
|
||||
vocab_size=129280,
|
||||
num_hidden_layers=43,
|
||||
num_attention_heads=32,
|
||||
num_key_value_heads=8,
|
||||
head_dim=128,
|
||||
intermediate_size=int(args.intermediate_size),
|
||||
max_position_embeddings=4096,
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
cfg = load_config(args.config)
|
||||
seed_all(int(cfg.seed))
|
||||
_init_single_process_group()
|
||||
|
||||
device = torch.device(
|
||||
args.device if (args.device != "cuda" or torch.cuda.is_available()) else "cpu"
|
||||
)
|
||||
precision_dtype = torch.bfloat16
|
||||
|
||||
model_args = cfg.model
|
||||
# Apply the documented smoke-only memory reductions.
|
||||
model_args.num_draft_layers = int(args.num_draft_layers)
|
||||
model_args.num_anchors = int(args.num_anchors)
|
||||
|
||||
target_config = _build_target_config(cfg, args)
|
||||
draft_config = build_draft_config(target_config=target_config, model_args=model_args)
|
||||
model = Qwen3DSparkModel(draft_config).to(device=device, dtype=precision_dtype)
|
||||
# build_models copies embed_tokens/lm_head from the target then freezes them.
|
||||
# We lack the target checkpoint, so we keep the random init and just freeze
|
||||
# (identical trainability; pipeline-only smoke).
|
||||
model.set_embedding_head_trainable(False)
|
||||
model.train()
|
||||
|
||||
dataset = CacheDataset(cache_dir=args.cache_dir)
|
||||
validate_train_cache(
|
||||
train_dataset=dataset,
|
||||
draft_model=model,
|
||||
target_model_name_or_path=model_args.target_model_name_or_path,
|
||||
)
|
||||
dataloader = DataLoader(
|
||||
dataset,
|
||||
batch_size=int(cfg.train.local_batch_size),
|
||||
shuffle=True,
|
||||
collate_fn=CacheCollator(),
|
||||
drop_last=True,
|
||||
)
|
||||
|
||||
optimizer = BF16Optimizer(
|
||||
model,
|
||||
lr=float(cfg.train.lr),
|
||||
total_steps=int(args.steps),
|
||||
warmup_ratio=float(cfg.train.warmup_ratio),
|
||||
weight_decay=float(cfg.train.weight_decay),
|
||||
)
|
||||
|
||||
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
|
||||
frozen = sum(p.numel() for p in model.parameters() if not p.requires_grad)
|
||||
print(
|
||||
f"[smoke] device={device} draft_layers={model_args.num_draft_layers} "
|
||||
f"num_anchors={model_args.num_anchors} intermediate={args.intermediate_size}\n"
|
||||
f"[smoke] params: trainable={trainable/1e6:.1f}M frozen(embed+head)="
|
||||
f"{frozen/1e6:.1f}M | dataset={len(dataset)} samples",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
data_iter = iter(dataloader)
|
||||
t0 = time.time()
|
||||
for step in range(int(args.steps)):
|
||||
try:
|
||||
batch = next(data_iter)
|
||||
except StopIteration:
|
||||
data_iter = iter(dataloader)
|
||||
batch = next(data_iter)
|
||||
batch = {
|
||||
key: value.to(device, non_blocking=True) for key, value in batch.items()
|
||||
}
|
||||
# Same boundary cast as the real path (cuda_prefetcher.move_batch_to_device):
|
||||
# embedding lookup requires int64, cache stores int32.
|
||||
if batch["input_ids"].dtype != torch.long:
|
||||
batch["input_ids"] = batch["input_ids"].to(torch.long)
|
||||
outputs = model(
|
||||
input_ids=batch["input_ids"],
|
||||
target_hidden_states=batch["target_hidden_states"],
|
||||
loss_mask=batch["loss_mask"],
|
||||
target_last_hidden_states=batch["target_last_hidden_states"],
|
||||
)
|
||||
loss = compute_dspark_loss(
|
||||
outputs=outputs,
|
||||
loss_decay_gamma=cfg.model.loss_decay_gamma,
|
||||
ce_loss_alpha=float(cfg.model.ce_loss_alpha),
|
||||
l1_loss_alpha=float(cfg.model.l1_loss_alpha),
|
||||
confidence_head_alpha=float(cfg.model.confidence_head_alpha),
|
||||
)
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
if device.type == "cuda":
|
||||
torch.cuda.synchronize()
|
||||
peak = (
|
||||
torch.cuda.max_memory_allocated() / 1024**3
|
||||
if device.type == "cuda"
|
||||
else 0.0
|
||||
)
|
||||
print(
|
||||
f"[smoke] step {step + 1}/{args.steps} loss={loss.item():.5f} "
|
||||
f"lr={optimizer.get_learning_rate():.3e} "
|
||||
f"peak_vram={peak:.2f}GB dt={time.time() - t0:.1f}s",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
dataset.close()
|
||||
if dist.is_initialized():
|
||||
dist.destroy_process_group()
|
||||
print("[smoke] OK: dataloader -> forward -> loss -> backward -> step verified.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user