File size: 8,360 Bytes
31dc8dc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 | """Thin launcher: EAGLE3 training through the SpecForge DataFlow runtime.
This script is a *launcher* (M3): it builds models + optimizer, hands them to the
runtime, and runs ``TrainerController.fit``. No training logic lives here — the
loop, loss, projection, checkpoint, and eval all live in ``specforge.runtime``.
Reuses the existing model/data builders from ``scripts.train_eagle3`` so model
construction stays DRY; only the *orchestration* moves behind the runtime. Both
modes converge at ``SampleRef`` and share one trainer/strategy/FSDP path:
* **offline** (``--train-hidden-states-path`` set): an ``OfflineManifestReader``
turns precomputed ``.ckpt`` files into refs.
* **online** (no hidden-states path): a ``RolloutWorker`` generates features from
the target model and commits refs onto the control plane's queue.
Examples:
# offline
torchrun --standalone --nproc_per_node 1 scripts/train_eagle3_dataflow.py \
--target-model-path <hf-model> --draft-model-config <cfg.json> \
--train-data-path <prompts.jsonl> --train-hidden-states-path <features_dir> \
--output-dir ./output --max-num-steps 20
# online (no --train-hidden-states-path)
torchrun --standalone --nproc_per_node 1 scripts/train_eagle3_dataflow.py \
--target-model-path <hf-model> --draft-model-config <cfg.json> \
--train-data-path <prompts.jsonl> --output-dir ./output --max-num-steps 20
"""
from accelerate.utils import set_seed
# reuse existing builders so model construction is not duplicated
from train_eagle3 import (
build_dataloaders,
build_draft_model,
build_target_model,
parse_args,
)
from specforge.distributed import destroy_distributed, init_distributed
from specforge.optimizer import BF16Optimizer
def _target_hidden_and_vocab(target_model):
"""Best-effort (hidden_size, vocab_size) from an Eagle3 target backend."""
cfg = getattr(getattr(target_model, "model", None), "config", None)
if cfg is not None:
return int(cfg.hidden_size), int(cfg.vocab_size)
raise RuntimeError(
"could not read hidden_size/vocab_size from the target model; pass them explicitly"
)
def _extract_prompts(train_dataloader):
"""Flatten the online train dataloader into metadata-only PromptTask payloads.
Each prompt carries only ``input_ids`` + ``loss_mask`` (the control plane
never holds tensors); ``attention_mask`` recovers the true unpadded length.
"""
prompts = []
for batch in train_dataloader:
input_ids = batch["input_ids"]
loss_mask = batch["loss_mask"]
attn = batch.get("attention_mask")
for i in range(input_ids.shape[0]):
n = int(attn[i].sum().item()) if attn is not None else input_ids.shape[1]
prompts.append(
{
"payload": {
"input_ids": input_ids[i, :n].tolist(),
"loss_mask": loss_mask[i, :n].tolist(),
}
}
)
return prompts
def main():
parser, args = parse_args()
# parse_args() does not derive target_batch_size (train_eagle3.main computes
# it inline before building dataloaders); the runtime builder and
# build_dataloaders both read it, so derive it here too.
args.target_batch_size = args.tp_size * args.batch_size
# TODO(dataflow-launcher parity with scripts/train_eagle3.py): this launcher
# covers core EAGLE3 training (offline + online: loss / projection / FSDP /
# TP / grad-accum / checkpoint), validated old-vs-new. The following
# train_eagle3.py features are NOT yet wired here and still require the
# legacy script:
# - VLM / multimodal targets (--is-vlm, QwenVLOnlineEagle3Model)
# - USP sequence parallelism (--attention-backend usp -> process_data_usp;
# this path uses OfflineEagle3Dataset.process_data, no per-rank seq shard)
# - eval loop (--eval-data-path / --eval-hidden-states-path)
# - resume from checkpoint (--resume)
# - experiment trackers (--report-to wandb / swanlab / tensorboard)
# - online multi-epoch re-rollout (online runs a single consume-once pass)
set_seed(args.seed)
init_distributed(
timeout=args.dist_timeout,
tp_size=args.tp_size,
sp_ring_size=args.sp_ring_size,
sp_ulysses_size=args.sp_ulysses_size,
)
online = args.train_hidden_states_path is None
draft_config, draft_model, _ckpt, _resume = build_draft_model(args)
# vocab mapping is produced from the prompt dataset exactly as today
train_dataloader, vocab_mapping_path, _eval = build_dataloaders(args, draft_config)
draft_model.load_vocab_mapping(vocab_mapping_path)
from specforge import OnlineEagle3Model
eagle3_model = OnlineEagle3Model(
draft_model=draft_model,
length=args.ttt_length,
attention_backend=args.attention_backend,
lk_loss_type=args.lk_loss_type,
kl_scale=args.kl_scale,
kl_decay=args.kl_decay,
).cuda()
# optimizer is built AFTER FSDP-wrap (inside the runtime) over the inner draft
def optimizer_factory(draft_module):
return BF16Optimizer(
draft_module,
lr=args.learning_rate,
max_grad_norm=args.max_grad_norm,
warmup_ratio=args.warmup_ratio,
total_steps=args.total_steps or 10_000,
)
logger = lambda m, s: print(f"step {s}: {m}", flush=True)
if online:
from specforge.runtime.launch import build_online_eagle3_runtime
# Online target produces features in-loop (any backend exposing
# generate_eagle3_data — HF or SGLang). is_online=True returns the model.
target_model, _ = build_target_model(args, draft_config, is_online=True)
hidden_size, vocab_size = _target_hidden_and_vocab(target_model)
prompts = _extract_prompts(train_dataloader)
print(f"[online] ingesting {len(prompts)} prompts for rollout", flush=True)
# num_epochs=1: the rollout output is a consume-once stream. Multi-epoch
# online (re-rollout each epoch) is a follow-up; one rollout pass here.
trainer, loader, workers, controller, drive_rollout = (
build_online_eagle3_runtime(
target_model=target_model,
prompts=prompts,
eagle3_model=eagle3_model,
optimizer_factory=optimizer_factory,
run_id="eagle3-online",
output_dir=args.output_dir,
target_hidden_size=hidden_size,
target_vocab_size=vocab_size,
target_repr="logits",
ttt_length=args.ttt_length,
batch_size=args.target_batch_size,
accumulation_steps=args.draft_accumulation_steps,
num_epochs=1,
max_steps=args.max_num_steps,
save_interval=args.save_interval,
tp_size=args.tp_size,
sp_ulysses_size=args.sp_ulysses_size,
sp_ring_size=args.sp_ring_size,
logger=logger,
)
)
produced = drive_rollout()
print(f"[online] rollout produced {produced} samples", flush=True)
trainer.fit(loader)
else:
from specforge.runtime.launch import build_offline_eagle3_runtime
target_head, _ = build_target_model(args, draft_config, is_online=False)
trainer, loader = build_offline_eagle3_runtime(
hidden_states_path=args.train_hidden_states_path,
eagle3_model=eagle3_model,
target_head=target_head,
optimizer_factory=optimizer_factory,
run_id="eagle3-offline",
output_dir=args.output_dir,
ttt_length=args.ttt_length,
max_len=args.max_length,
batch_size=args.target_batch_size,
accumulation_steps=args.draft_accumulation_steps,
num_epochs=args.num_epochs,
max_steps=args.max_num_steps,
save_interval=args.save_interval,
tp_size=args.tp_size,
sp_ulysses_size=args.sp_ulysses_size,
sp_ring_size=args.sp_ring_size,
logger=logger,
)
trainer.fit(loader)
destroy_distributed()
if __name__ == "__main__":
main()
|