1435 lines
80 KiB
C++
1435 lines
80 KiB
C++
// ace-train.cpp — HOT-Step training toolchain.
|
|
// Phase 2 ships one subcommand: `preprocess`.
|
|
//
|
|
// Standalone tool (no acestep-core): every engine module it uses is a
|
|
// header-only `static` API. Machine-readable JSONL on stdout (--jsonl),
|
|
// human logs on stderr. Exit: 0 ok, 1 runtime failure, 2 usage.
|
|
//
|
|
// docs/plans/2026-07-27-preprocess-implementation.md §3.1
|
|
|
|
#include "model-registry.h"
|
|
#include "train/dit-train-run.h" // pulls in every dit-*.h (DiT LoRA trainer)
|
|
#include "train/lm-train-run.h" // pulls in every lm-*.h (LM LoRA trainer)
|
|
#include "train/preprocess-run.h" // pulls in st-write.h + preprocess-io.h
|
|
#include "train/spike-run.h" // Phase-0 training spike (evidence code)
|
|
|
|
#include <cstdarg>
|
|
#include <cstdio>
|
|
#include <cstdlib>
|
|
#include <cstring>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
#ifdef _WIN32
|
|
# include <windows.h>
|
|
#else
|
|
# include <dirent.h>
|
|
#endif
|
|
|
|
// ─── JSONL emitter ──────────────────────────────────────────────────────────
|
|
|
|
static bool g_jsonl = false;
|
|
|
|
static void jl(const char * fmt, ...) { // printf-style, caller writes the JSON
|
|
if (!g_jsonl) {
|
|
return;
|
|
}
|
|
va_list ap;
|
|
va_start(ap, fmt);
|
|
vprintf(fmt, ap);
|
|
va_end(ap);
|
|
putchar('\n');
|
|
fflush(stdout); // MANDATORY: Node reads line-by-line
|
|
}
|
|
|
|
static std::string json_escape(const std::string & s) {
|
|
std::string o;
|
|
o.reserve(s.size() + 8);
|
|
for (size_t i = 0; i < s.size(); i++) {
|
|
unsigned char c = (unsigned char) s[i];
|
|
switch (c) {
|
|
case '"':
|
|
o += "\\\"";
|
|
break;
|
|
case '\\':
|
|
o += "\\\\";
|
|
break;
|
|
case '\b':
|
|
o += "\\b";
|
|
break;
|
|
case '\f':
|
|
o += "\\f";
|
|
break;
|
|
case '\n':
|
|
o += "\\n";
|
|
break;
|
|
case '\r':
|
|
o += "\\r";
|
|
break;
|
|
case '\t':
|
|
o += "\\t";
|
|
break;
|
|
default:
|
|
if (c < 0x20) {
|
|
char b[8];
|
|
snprintf(b, sizeof(b), "\\u%04x", c);
|
|
o += b;
|
|
} else {
|
|
o.push_back((char) c);
|
|
}
|
|
}
|
|
}
|
|
return o;
|
|
}
|
|
|
|
static void jl_stage(const char * stage, const char * path, long long ms) {
|
|
jl("{\"type\":\"model\",\"stage\":\"%s\",\"path\":\"%s\",\"ms\":%lld}", stage, json_escape(path).c_str(), ms);
|
|
fprintf(stderr, "[ace-train] loaded %s (%lld ms)\n", stage, ms);
|
|
}
|
|
|
|
static void jl_warn(const char * msg) {
|
|
jl("{\"type\":\"log\",\"level\":\"warn\",\"message\":\"%s\"}", json_escape(msg).c_str());
|
|
fprintf(stderr, "[ace-train] WARN: %s\n", msg);
|
|
}
|
|
|
|
// ─── usage ──────────────────────────────────────────────────────────────────
|
|
|
|
static void print_usage(void) {
|
|
fprintf(stderr,
|
|
"ace-train %s — HOT-Step training toolchain\n"
|
|
"\n"
|
|
"usage: ace-train <subcommand> [options]\n"
|
|
"\n"
|
|
"Subcommands:\n"
|
|
" preprocess Build per-song tensor caches from a dataset.\n"
|
|
" train-lm Train a planner-LM LoRA from a preprocessed tensor cache.\n"
|
|
" train-dit Train a DiT LoRA from a preprocessed tensor cache.\n"
|
|
"\n"
|
|
"ace-train preprocess (all paths absolute; long options only; \"--flag value\" form)\n"
|
|
"\n"
|
|
" Input (exactly one required):\n"
|
|
" --manifest <path> preprocess_manifest.json (server path)\n"
|
|
" --dataset <path> Side-Step dataset.json (standalone path)\n"
|
|
"\n"
|
|
" Output (required):\n"
|
|
" --out <dir> output directory; created if missing\n"
|
|
"\n"
|
|
" Models (required):\n"
|
|
" --models <dir> root scanned for name lookups\n"
|
|
" --dit <name|path> DiT GGUF/dir — cond-enc weights, silence_latent, variant identity\n"
|
|
" --vae <name|path> VAE GGUF/safetensors (must NOT be .onnx)\n"
|
|
" --text-enc <name|path> Qwen3-Embedding GGUF/dir (also supplies BPE vocab)\n"
|
|
"\n"
|
|
" Options (defaults shown):\n"
|
|
" --max-duration <sec> 240 0 = no truncation\n"
|
|
" --normalize <none|peak> peak\n"
|
|
" --target-db <db> -1.0 peak-normalize target\n"
|
|
" --dtype <f32|bf16> f32 storage dtype for every tensor\n"
|
|
" --compat <hotstep|sidestep> hotstep\n"
|
|
" --timbre <silence|zeros> silence\n"
|
|
" --max-caption-tokens <n> 256\n"
|
|
" --max-lyric-tokens <n> 512\n"
|
|
" --vae-chunk <latent-frames> 384\n"
|
|
" --vae-overlap <latent-frames> 48\n"
|
|
" --ffmpeg <path> (none) required for non-WAV/MP3 input\n"
|
|
" --overwrite off re-encode songs that already have a cache\n"
|
|
" --limit <n> 0 debug: stop after n songs (0 = all)\n"
|
|
" --jsonl off emit machine-readable JSONL on stdout\n"
|
|
" --verify <file> debug: re-open a written .safetensors and report it\n"
|
|
" -h, --help\n"
|
|
"\n"
|
|
"ace-train train-lm (all paths absolute; long options only; \"--flag value\" form)\n"
|
|
"\n"
|
|
" Stages (default: extract,train,export):\n"
|
|
" --stages <csv> subset of extract,train,export in that order\n"
|
|
"\n"
|
|
" Input:\n"
|
|
" --tensors <dir> preprocess variant dir (required for `extract`;\n"
|
|
" also the default location of the codes file)\n"
|
|
" --codes <path> lm_codes.jsonl (default <tensors>/lm_codes.jsonl)\n"
|
|
" --out <dir> adapter output dir (required for train/export)\n"
|
|
"\n"
|
|
" Models:\n"
|
|
" --models <dir> root scanned by registry_scan() for name lookups\n"
|
|
" --dit <name|path> FSQ-tokenizer source for `extract`\n"
|
|
" (default: dit_path from <tensors>/preprocess_meta.json)\n"
|
|
" --lm <name|path> LM base GGUF (required for `train`)\n"
|
|
" --lm-size <0.6B|1.7B|4B> label written into adapter_config/base_model_name_or_path\n"
|
|
" (default: inferred from the LM's hidden_size/n_layers)\n"
|
|
"\n"
|
|
" LoRA:\n"
|
|
" --rank <n> 16\n"
|
|
" --alpha <n> 32\n"
|
|
"\n"
|
|
" Optimizer / schedule:\n"
|
|
" --lr <f> 1e-4\n"
|
|
" --epochs <n> 16 hard cap; target-loss usually stops earlier\n"
|
|
" --grad-accum <n> 4\n"
|
|
" --warmup-ratio <f> 0.05\n"
|
|
" --grad-clip <f> 1.0 0 = disabled\n"
|
|
" --weight-decay <f> 0.01\n"
|
|
" --seed <n> 42\n"
|
|
" --target-loss <f> 0.4 0 = disabled (run to the epoch cap)\n"
|
|
" --order <shuffle|fixed> shuffle `fixed` = file order (A/B parity runs)\n"
|
|
"\n"
|
|
" Sequence / memory:\n"
|
|
" --max-len <n> 0 0 = auto-fit from free VRAM\n"
|
|
" --vram-reserve-mb <n> 1024 headroom left unallocated\n"
|
|
" --low-vram <auto|on|off> auto per-layer checkpointing + chunked CE;\n"
|
|
" auto = ON iff --lm-size is 4B, or the naive\n"
|
|
" auto-fit would yield max-len < 2048\n"
|
|
" --attn-head-block <n> -1 heads per attention block inside a segment;\n"
|
|
" 0 = whole-head attention (naive behaviour),\n"
|
|
" -1 = engine picks (<=16 heads -> 0, else 8).\n"
|
|
" Must divide n_heads AND n*n_kv/n_heads >= 1\n"
|
|
" --lm-chunk <n> 128 trained positions per CE chunk (low-vram only)\n"
|
|
" --weights <f32-window|bf16> f32-window projection GEMM dtype.\n"
|
|
" f32-window = the shipped per-segment F32 weight\n"
|
|
" cast (cublasSgemm/TF32).\n"
|
|
" bf16 = BF16 projections + backward surgery\n"
|
|
" (cublasGemmEx BF16). Needs a BF16\n"
|
|
" base and low-VRAM mode, and CHANGES\n"
|
|
" THE TRAINED WEIGHTS (the activation\n"
|
|
" gradient is BF16-rounded at every\n"
|
|
" layer). Not resume-compatible with\n"
|
|
" an f32-window run.\n"
|
|
" EXPERIMENTAL, NOT ACCEPTED: its own parity gates\n"
|
|
" (--self-test T14/T15) measure OVER their bars —\n"
|
|
" max rel ~1.2e-1 vs a 2e-2 bar, cosine ~0.9995 vs\n"
|
|
" 1-1e-4. Direction is close but not within spec;\n"
|
|
" judge a bf16 adapter by ear before trusting it.\n"
|
|
" --bwd <outprod|mm> outprod MUL_MAT activation-gradient formulation.\n"
|
|
" outprod = upstream ggml out_prod (F32-only on\n"
|
|
" CUDA; forces an F32 weight, so the\n"
|
|
" forward runs TF32).\n"
|
|
" mm = mul_mat(cont(transpose(W)), grad) —\n"
|
|
" identical maths, dtype-agnostic, so a\n"
|
|
" BF16 weight uses BF16 tensor cores.\n"
|
|
" ~1.7-1.8x per layer per step on an\n"
|
|
" RTX 5090. Needs the vendored patch\n"
|
|
" engine/patches/mm-backward.patch.\n"
|
|
" --optimizer <adamw|muon> adamw muon puts every 2-D parameter whose SHORT side is\n"
|
|
" >= --muon-min-dim on orthogonalized-momentum\n"
|
|
" (Newton-Schulz) updates. FOR A LoRA THE SHORT\n"
|
|
" SIDE IS THE RANK, so at rank 16 (the default)\n"
|
|
" every adapter matrix qualifies and at rank 8\n"
|
|
" none do — the startup line reports the split.\n"
|
|
" On the DiT this measured 1.41x fewer epochs to\n"
|
|
" target; it is untested on the LM.\n"
|
|
" --muon-lr-scale <f> 1.0 multiplies the shared LR schedule for Muon\n"
|
|
" params only. Muon's update is normalized, so its\n"
|
|
" LR does NOT mean AdamW's — the DiT needed ~20.\n"
|
|
" --muon-momentum <f> 0.95\n"
|
|
" --muon-ns-steps <n> 5 Newton-Schulz iterations\n"
|
|
" --muon-min-dim <n> 16 short-side floor for Muon eligibility\n"
|
|
" --muon-bucket <n> 16 params per batched Newton-Schulz (measured\n"
|
|
" optimum 8-16; larger is worse, the gather is\n"
|
|
" quadratic in bucket size)\n"
|
|
" --no-muon-nesterov use the plain momentum buffer, not g + mu*m\n"
|
|
" --batch <n|auto> 1 micro-batch size. NOT SUPPORTED in this build:\n"
|
|
" micro-batching was never written — the host\n"
|
|
" overhead it would amortise measures 9.3%% at 4B,\n"
|
|
" under the 10%% bar its build gate required.\n"
|
|
" Use --grad-accum to change effective batch size.\n"
|
|
" --loss-on-cot default ON\n"
|
|
" --no-loss-on-cot\n"
|
|
"\n"
|
|
" Adapter identity:\n"
|
|
" --trigger <word> \"\" trigger word embedded in the adapter's\n"
|
|
" metadata. Default: custom_tag from the\n"
|
|
" variant's preprocess_meta.json.\n"
|
|
" --trigger-position <prepend|append> where the trigger sat in the training\n"
|
|
" captions. Default: that file's tag_position.\n"
|
|
"\n"
|
|
" Run management:\n"
|
|
" --milestone-step <f> 0.1 0 = disabled\n"
|
|
" --milestone-keep <n> 6\n"
|
|
" --no-milestones\n"
|
|
" --overwrite off re-extract every song / overwrite <out>\n"
|
|
" --limit <n> 0 debug: first n songs only\n"
|
|
" --self-test off run the correctness gates and exit\n"
|
|
" --jsonl off machine-readable JSONL on stdout\n"
|
|
" -h, --help\n"
|
|
"\n"
|
|
"ace-train train-dit (all paths absolute; long options only; \"--flag value\" form)\n"
|
|
"\n"
|
|
" Stages (default: train,export):\n"
|
|
" --stages <csv> subset of train,export in that order\n"
|
|
"\n"
|
|
" Input / output:\n"
|
|
" --tensors <dir> preprocess variant dir (required)\n"
|
|
" --out <dir> adapter output dir (required)\n"
|
|
" --models <dir> root scanned by registry_scan() for name lookups\n"
|
|
" --dit <name|path> base DiT GGUF (default: dit_path from\n"
|
|
" <tensors>/preprocess_meta.json — must match the\n"
|
|
" variant the latents were made against)\n"
|
|
"\n"
|
|
" Adapter:\n"
|
|
" --adapter-type <lora|lokr> lora\n"
|
|
" --rank <n> 128 lora only\n"
|
|
" --alpha <n> 256 lora only\n"
|
|
" --lokr-dim <n> 512 lokr only, 4-4096\n"
|
|
" --lokr-alpha <f> 512 lokr only, 0-8192; 0 = dim. LyCORIS\n"
|
|
" forces alpha = dim (scale 1) wherever\n"
|
|
" both kron factors are monolithic.\n"
|
|
" --lokr-factor <n> 6 lokr only, -1 or 2-64\n"
|
|
" --lokr-decompose-both / --no-lokr-decompose-both default ON (parity knob;\n"
|
|
" inert at dim 512 — w1 is always monolithic)\n"
|
|
" --target-mlp / --no-target-mlp default ON also LoRA mlp gate/up/down\n"
|
|
" --layers <n> 0 0 = auto; else train the top n layers\n"
|
|
"\n"
|
|
" Objective (design 4.5):\n"
|
|
" --loss-weighting <none|flow_snr> flow_snr\n"
|
|
" --snr-gamma <f> 5.0\n"
|
|
" --t-bias <f> 0.5\n"
|
|
" --channel-balance / --no-channel-balance default ON\n"
|
|
" --timestep-mu <f> -0.4\n"
|
|
" --timestep-sigma <f> 1.0\n"
|
|
" --t-min <f> 0.0 interval-expert window (rejection resample)\n"
|
|
" --t-max <f> 1.0\n"
|
|
" --cfg-ratio <f> 0.15\n"
|
|
" --genre-ratio <n> 30 percent, 0-100\n"
|
|
"\n"
|
|
" Optimizer / schedule:\n"
|
|
" --lr <f> 5e-4\n"
|
|
" --epochs <n> 400 hard cap; target-loss usually stops earlier\n"
|
|
" --batch <n> 1 1-16; crops per micro-batch, from that many\n"
|
|
" DIFFERENT songs. Reduced with a warn when the\n"
|
|
" variant has fewer songs, or when\n"
|
|
" n_kv_heads*max(S,enc_S)*B would exceed ggml's\n"
|
|
" CUDA repeat_back cap. Default 1 (off): measured\n"
|
|
" ~2.5x SLOWER at full depth on a 32 GB card, but\n"
|
|
" ~2.4x FASTER on shallow/partial-depth runs.\n"
|
|
" --grad-accum <n> 4 counts MICRO-BATCHES, so an optimizer step sees\n"
|
|
" batch x grad-accum songs (4 at the defaults).\n"
|
|
" --ckpt <n> 1 gradient checkpointing: 0 = off (one monolithic\n"
|
|
" fwd+bwd graph), 1 = auto (the VRAM fit picks the\n"
|
|
" segment count), 2-32 = exactly that many segments.\n"
|
|
" Segments trade one extra no-grad forward (~+30-40%%\n"
|
|
" of forward time) for an activation set divided by\n"
|
|
" the segment count. Clamped to the trained depth.\n"
|
|
" --warmup-ratio <f> 0.05\n"
|
|
" --grad-clip <f> 1.0 0 = disabled\n"
|
|
" --weight-decay <f> 0.01\n"
|
|
" --seed <n> 42\n"
|
|
" --target-loss <f> 0.4 0 = disabled (run to the epoch cap)\n"
|
|
" --order <shuffle|fixed> shuffle\n"
|
|
"\n"
|
|
" Crop / memory:\n"
|
|
" --crop <n> 0 latent frames; 0 = auto-fit\n"
|
|
" --crop-min <n> 375\n"
|
|
" --crop-max <n> 1250\n"
|
|
" --vram-reserve-mb <n> 2048 desktop/OS headroom left unallocated\n"
|
|
" --vram-safety <f> 0.05 extra margin on the footprint model\n"
|
|
" (0.12 for --adapter-type lokr unless set)\n"
|
|
" --mirror <f32|bf16> f32 frozen-weight mirror precision. bf16 keeps\n"
|
|
" the trainable layers' matmul weights in the\n"
|
|
" base's native BF16 instead of promoting them\n"
|
|
" to F32, roughly halving the mirror. CUDA only\n"
|
|
" (engine/patches/bf16-out-prod.patch);\n"
|
|
" EXPERIMENTAL, adapter quality unvalidated.\n"
|
|
" --bwd <outprod|mm> outprod MUL_MAT activation-gradient formulation.\n"
|
|
" outprod = upstream ggml out_prod (F32-only on\n"
|
|
" CUDA; forces an F32 weight, so the\n"
|
|
" forward runs TF32).\n"
|
|
" mm = mul_mat(cont(transpose(W)), grad) —\n"
|
|
" identical maths, dtype-agnostic, so a\n"
|
|
" BF16 mirror uses BF16 tensor cores.\n"
|
|
" ~1.7-1.8x per layer per step on an\n"
|
|
" RTX 5090. Needs the vendored patch\n"
|
|
" engine/patches/mm-backward.patch.\n"
|
|
"\n"
|
|
" Optimizer:\n"
|
|
" --optimizer <adamw|muon> adamw muon puts every 2-D parameter whose SHORT side is\n"
|
|
" >= --muon-min-dim on orthogonalized-momentum\n"
|
|
" (Newton-Schulz) updates, and leaves the rest on\n"
|
|
" AdamW — a LoKR w1 is [4,5], where orthogonalizing\n"
|
|
" means nothing. Muon needs only ONE momentum\n"
|
|
" buffer where AdamW needs two (~900 MB at dim512).\n"
|
|
" --muon-lr-scale <f> 1.0 multiplies the shared LR schedule, for Muon\n"
|
|
" parameters only. Muon's update is normalized by\n"
|
|
" construction, so its LR does NOT mean AdamW's —\n"
|
|
" expect to retune before comparing.\n"
|
|
" --muon-momentum <f> 0.95\n"
|
|
" --muon-ns-steps <n> 5 Newton-Schulz iterations\n"
|
|
" --muon-min-dim <n> 16 short-side floor for Muon eligibility\n"
|
|
" --no-muon-nesterov use the plain momentum buffer instead of g + mu*m\n"
|
|
"\n"
|
|
" Diagnostics:\n"
|
|
" --profile-step <n> 0 0 = off. n > 0 times every micro-step into\n"
|
|
" assemble / upload / build / backward / alloc /\n"
|
|
" compute / readback / free and prints the mean\n"
|
|
" every n steps, with graph node, leaf and\n"
|
|
" scheduler split counts. The run-mean lands in\n"
|
|
" dit_train_log.json under runtime.profile_ms.\n"
|
|
" Needs --ckpt 0 (the segmented path is not\n"
|
|
" instrumented). Off, the graph compute path is\n"
|
|
" byte-identical to a build without this flag.\n"
|
|
" --profile-ops off times ONE warmed-up micro-step node by node and\n"
|
|
" prints the top op/shape keys by share. The eval\n"
|
|
" callback serialises the graph, so absolute times\n"
|
|
" are inflated - read the shares, not the ms.\n"
|
|
"\n"
|
|
" Adapter identity:\n"
|
|
" --trigger <word> \"\" trigger word embedded in the adapter's\n"
|
|
" metadata. Default: custom_tag from the\n"
|
|
" variant's preprocess_meta.json.\n"
|
|
" --trigger-position <prepend|append> where the trigger sat in the training\n"
|
|
" captions. Default: that file's tag_position.\n"
|
|
"\n"
|
|
" Run management:\n"
|
|
" --milestone-step <f> 0.1 0 = disabled\n"
|
|
" --milestone-keep <n> 6\n"
|
|
" --no-milestones\n"
|
|
" --overwrite off\n"
|
|
" --limit <n> 0 debug: first n songs only\n"
|
|
" --self-test off run the correctness gates and exit\n"
|
|
" --jsonl off machine-readable JSONL on stdout\n"
|
|
" -h, --help\n",
|
|
ACE_VERSION);
|
|
}
|
|
|
|
// ─── path helpers ───────────────────────────────────────────────────────────
|
|
|
|
static bool looks_like_path(const std::string & v) {
|
|
return v.find('/') != std::string::npos || v.find('\\') != std::string::npos || pm_path_exists(v);
|
|
}
|
|
|
|
// Resolve a model argument: explicit path, else exact `name` lookup in a registry bucket.
|
|
static bool resolve_model(const std::string & arg, const std::vector<ModelEntry> & bucket, bool non_onnx,
|
|
std::string & out, std::string & name_out) {
|
|
if (arg.empty()) {
|
|
return false;
|
|
}
|
|
if (looks_like_path(arg)) {
|
|
if (!pm_path_exists(arg)) {
|
|
return false;
|
|
}
|
|
// The registry branch below filters ONNX out via registry_find_non_onnx;
|
|
// an explicit path must obey the same rule or a decoder-only ONNX VAE
|
|
// reaches vae_enc_load(), which fatal-exits with no diagnostic.
|
|
if (non_onnx && arg.size() >= 5 && pm_lower(arg.substr(arg.size() - 5)) == ".onnx") {
|
|
return false;
|
|
}
|
|
out = arg;
|
|
name_out = pm_basename(arg);
|
|
return true;
|
|
}
|
|
const ModelEntry * e = non_onnx ? registry_find_non_onnx(bucket, arg.c_str()) : registry_find(bucket, arg.c_str());
|
|
if (!e) {
|
|
return false;
|
|
}
|
|
out = e->path;
|
|
name_out = e->name;
|
|
return true;
|
|
}
|
|
|
|
// Delete every file in `dir` whose name ends with `suffix` (non-recursive).
|
|
static void purge_suffix(const std::string & dir, const char * suffix) {
|
|
std::vector<std::string> names;
|
|
registry_list_dir(dir.c_str(), &names);
|
|
const size_t slen = strlen(suffix);
|
|
for (size_t i = 0; i < names.size(); i++) {
|
|
if (names[i].size() >= slen && names[i].compare(names[i].size() - slen, slen, suffix) == 0) {
|
|
remove((dir + "/" + names[i]).c_str());
|
|
}
|
|
}
|
|
}
|
|
|
|
static void purge_all(const std::string & dir) {
|
|
std::vector<std::string> names;
|
|
registry_list_dir(dir.c_str(), &names);
|
|
for (size_t i = 0; i < names.size(); i++) {
|
|
remove((dir + "/" + names[i]).c_str());
|
|
}
|
|
}
|
|
|
|
// ─── --verify (debug round-trip reader, §6.2 assertion 6) ───────────────────
|
|
|
|
static int cmd_verify(const std::string & file) {
|
|
STFile st;
|
|
if (!st_open(&st, file.c_str())) {
|
|
fprintf(stderr, "[verify] cannot open %s\n", file.c_str());
|
|
return 1;
|
|
}
|
|
bool ok = true;
|
|
fprintf(stderr, "[verify] %s — %zu tensors, data_offset=%zu, size=%zu\n", file.c_str(), st.entries.size(),
|
|
st.data_offset, st.file_size);
|
|
for (size_t i = 0; i < st.entries.size(); i++) {
|
|
const STEntry & e = st.entries[i];
|
|
long long n = 1;
|
|
for (int d = 0; d < e.n_dims; d++) {
|
|
n *= e.shape[d];
|
|
}
|
|
const size_t esz = (e.dtype == "F32") ? 4u : 2u;
|
|
const size_t want = (size_t) n * esz;
|
|
const size_t have = e.data_end - e.data_start;
|
|
const bool inrng = st.data_offset + e.data_end <= st.file_size;
|
|
bool fin = true;
|
|
double mn = 0, mx = 0, sum = 0;
|
|
const void * d = st_data(st, e);
|
|
if (d && inrng && want == have) {
|
|
for (long long k = 0; k < n; k++) {
|
|
float v = (e.dtype == "F32") ? ((const float *) d)[k] : pm_bf16_to_f32(((const uint16_t *) d)[k]);
|
|
if (!std::isfinite(v)) {
|
|
fin = false;
|
|
}
|
|
if (k == 0 || v < mn) {
|
|
mn = v;
|
|
}
|
|
if (k == 0 || v > mx) {
|
|
mx = v;
|
|
}
|
|
sum += v;
|
|
}
|
|
}
|
|
const bool entry_ok = d && inrng && want == have && fin;
|
|
ok = ok && entry_ok;
|
|
fprintf(stderr, "[verify] %-32s %-5s [", e.name.c_str(), e.dtype.c_str());
|
|
for (int d2 = 0; d2 < e.n_dims; d2++) {
|
|
fprintf(stderr, "%s%lld", d2 ? "," : "", (long long) e.shape[d2]);
|
|
}
|
|
fprintf(stderr, "] bytes=%zu/%zu finite=%s min=%.6g max=%.6g mean=%.6g %s\n", have, want, fin ? "yes" : "NO",
|
|
mn, mx, n ? sum / (double) n : 0.0, entry_ok ? "OK" : "FAIL");
|
|
if (g_jsonl) {
|
|
std::string shape;
|
|
for (int d2 = 0; d2 < e.n_dims; d2++) {
|
|
char b[32];
|
|
snprintf(b, sizeof(b), "%s%lld", d2 ? "," : "", (long long) e.shape[d2]);
|
|
shape += b;
|
|
}
|
|
jl("{\"type\":\"verify\",\"name\":\"%s\",\"dtype\":\"%s\",\"shape\":[%s],\"bytes\":%zu,\"expected\":%zu,"
|
|
"\"finite\":%s,\"min\":%.9g,\"max\":%.9g,\"mean\":%.9g,\"ok\":%s}",
|
|
json_escape(e.name).c_str(), e.dtype.c_str(), shape.c_str(), have, want, fin ? "true" : "false", mn, mx,
|
|
n ? sum / (double) n : 0.0, entry_ok ? "true" : "false");
|
|
}
|
|
}
|
|
st_close(&st);
|
|
return ok ? 0 : 1;
|
|
}
|
|
|
|
// ─── preprocess ─────────────────────────────────────────────────────────────
|
|
|
|
static int cmd_preprocess(int argc, char ** argv) {
|
|
PreprocessOpts o;
|
|
std::string manifest_path, dataset_path, models_dir;
|
|
std::string dit_arg, vae_arg, text_arg, verify_file;
|
|
std::string normalize = "peak", dtype = "f32", compat = "hotstep", timbre = "silence";
|
|
int limit = 0;
|
|
|
|
for (int i = 1; i < argc; i++) {
|
|
if (!strcmp(argv[i], "--manifest") && i + 1 < argc) manifest_path = argv[++i];
|
|
else if (!strcmp(argv[i], "--dataset") && i + 1 < argc) dataset_path = argv[++i];
|
|
else if (!strcmp(argv[i], "--out") && i + 1 < argc) o.out_dir = argv[++i];
|
|
else if (!strcmp(argv[i], "--models") && i + 1 < argc) models_dir = argv[++i];
|
|
else if (!strcmp(argv[i], "--dit") && i + 1 < argc) dit_arg = argv[++i];
|
|
else if (!strcmp(argv[i], "--vae") && i + 1 < argc) vae_arg = argv[++i];
|
|
else if (!strcmp(argv[i], "--text-enc") && i + 1 < argc) text_arg = argv[++i];
|
|
else if (!strcmp(argv[i], "--max-duration") && i + 1 < argc) o.max_duration = atoi(argv[++i]);
|
|
else if (!strcmp(argv[i], "--normalize") && i + 1 < argc) normalize = argv[++i];
|
|
else if (!strcmp(argv[i], "--target-db") && i + 1 < argc) o.target_db = (float) atof(argv[++i]);
|
|
else if (!strcmp(argv[i], "--dtype") && i + 1 < argc) dtype = argv[++i];
|
|
else if (!strcmp(argv[i], "--compat") && i + 1 < argc) compat = argv[++i];
|
|
else if (!strcmp(argv[i], "--timbre") && i + 1 < argc) timbre = argv[++i];
|
|
else if (!strcmp(argv[i], "--max-caption-tokens") && i + 1 < argc) o.max_caption_tokens = atoi(argv[++i]);
|
|
else if (!strcmp(argv[i], "--max-lyric-tokens") && i + 1 < argc) o.max_lyric_tokens = atoi(argv[++i]);
|
|
else if (!strcmp(argv[i], "--vae-chunk") && i + 1 < argc) o.vae_chunk = atoi(argv[++i]);
|
|
else if (!strcmp(argv[i], "--vae-overlap") && i + 1 < argc) o.vae_overlap = atoi(argv[++i]);
|
|
else if (!strcmp(argv[i], "--ffmpeg") && i + 1 < argc) o.ffmpeg = argv[++i];
|
|
else if (!strcmp(argv[i], "--limit") && i + 1 < argc) limit = atoi(argv[++i]);
|
|
else if (!strcmp(argv[i], "--verify") && i + 1 < argc) verify_file = argv[++i];
|
|
else if (!strcmp(argv[i], "--overwrite")) o.overwrite = true;
|
|
else if (!strcmp(argv[i], "--jsonl")) g_jsonl = true;
|
|
else if (!strcmp(argv[i], "-h") || !strcmp(argv[i], "--help")) { print_usage(); return 0; }
|
|
else {
|
|
fprintf(stderr, "ace-train: unknown option '%s'\n", argv[i]);
|
|
print_usage();
|
|
return 2;
|
|
}
|
|
}
|
|
|
|
if (!verify_file.empty()) {
|
|
return cmd_verify(verify_file);
|
|
}
|
|
|
|
// ── validate arguments ──────────────────────────────────────────────
|
|
if (manifest_path.empty() && dataset_path.empty()) {
|
|
fprintf(stderr, "ace-train preprocess: --manifest or --dataset is required\n");
|
|
return 2;
|
|
}
|
|
if (o.out_dir.empty()) {
|
|
fprintf(stderr, "ace-train preprocess: --out is required\n");
|
|
return 2;
|
|
}
|
|
if (models_dir.empty()) {
|
|
fprintf(stderr, "ace-train preprocess: --models is required\n");
|
|
return 2;
|
|
}
|
|
if (dit_arg.empty() || vae_arg.empty() || text_arg.empty()) {
|
|
fprintf(stderr, "ace-train preprocess: --dit, --vae and --text-enc are all required\n");
|
|
return 2;
|
|
}
|
|
if (normalize != "none" && normalize != "peak") {
|
|
fprintf(stderr, "ace-train preprocess: --normalize must be none|peak\n");
|
|
return 2;
|
|
}
|
|
if (dtype != "f32" && dtype != "bf16") {
|
|
fprintf(stderr, "ace-train preprocess: --dtype must be f32|bf16\n");
|
|
return 2;
|
|
}
|
|
if (compat != "hotstep" && compat != "sidestep") {
|
|
fprintf(stderr, "ace-train preprocess: --compat must be hotstep|sidestep\n");
|
|
return 2;
|
|
}
|
|
if (timbre != "silence" && timbre != "zeros") {
|
|
fprintf(stderr, "ace-train preprocess: --timbre must be silence|zeros\n");
|
|
return 2;
|
|
}
|
|
if (o.max_duration < 0 || o.vae_chunk < 64 || o.vae_overlap < 0 || o.vae_overlap >= o.vae_chunk ||
|
|
o.max_caption_tokens < 16 || o.max_lyric_tokens < 16) {
|
|
fprintf(stderr, "ace-train preprocess: numeric option out of range\n");
|
|
return 2;
|
|
}
|
|
|
|
o.normalize_peak = (normalize == "peak");
|
|
o.dtype = (dtype == "bf16") ? STW_BF16 : STW_F32;
|
|
o.sidestep_compat = (compat == "sidestep");
|
|
o.timbre_zeros = (timbre == "zeros");
|
|
|
|
// ── resolve models ──────────────────────────────────────────────────
|
|
ModelRegistry reg;
|
|
registry_scan(®, models_dir.c_str());
|
|
|
|
std::string dit_name, vae_name, text_name;
|
|
if (!resolve_model(dit_arg, reg.dit, false, o.dit_path, dit_name)) {
|
|
fprintf(stderr, "ace-train preprocess: cannot resolve --dit '%s' in %s\n", dit_arg.c_str(), models_dir.c_str());
|
|
return 2;
|
|
}
|
|
if (!resolve_model(vae_arg, reg.vae, true, o.vae_path, vae_name)) {
|
|
fprintf(stderr, "ace-train preprocess: cannot resolve --vae '%s' in %s (ONNX VAEs cannot encode)\n",
|
|
vae_arg.c_str(), models_dir.c_str());
|
|
return 2;
|
|
}
|
|
if (!resolve_model(text_arg, reg.text_enc, false, o.text_enc_path, text_name)) {
|
|
fprintf(stderr, "ace-train preprocess: cannot resolve --text-enc '%s' in %s\n", text_arg.c_str(),
|
|
models_dir.c_str());
|
|
return 2;
|
|
}
|
|
o.model_variant = dit_name;
|
|
|
|
// ── output dir ──────────────────────────────────────────────────────
|
|
if (!pm_mkdir_p(o.out_dir)) {
|
|
fprintf(stderr, "ace-train preprocess: cannot create output directory %s\n", o.out_dir.c_str());
|
|
return 1;
|
|
}
|
|
pm_mkdir_p(o.out_dir + "/.tmp");
|
|
purge_suffix(o.out_dir, ".__writing__");
|
|
purge_suffix(o.out_dir, ".__old__"); // aside copies from an interrupted replace
|
|
purge_all(o.out_dir + "/.tmp");
|
|
|
|
// ── load the manifest ───────────────────────────────────────────────
|
|
PManifest mf;
|
|
std::string err;
|
|
bool loaded = false;
|
|
if (!manifest_path.empty()) {
|
|
loaded = pm_load_manifest(manifest_path.c_str(), mf, err); // --manifest wins if both given
|
|
} else {
|
|
loaded = pm_load_dataset_json(dataset_path.c_str(), mf, err);
|
|
}
|
|
if (!loaded) {
|
|
fprintf(stderr, "ace-train preprocess: %s\n", err.c_str());
|
|
return 1;
|
|
}
|
|
|
|
int total = (int) mf.samples.size();
|
|
if (limit > 0 && limit < total) {
|
|
total = limit;
|
|
}
|
|
|
|
ggml_time_init();
|
|
const int64_t t_run0 = ggml_time_ms();
|
|
|
|
jl("{\"type\":\"start\",\"total\":%d,\"out\":\"%s\",\"variant\":\"%s\",\"compat\":\"%s\",\"dtype\":\"%s\","
|
|
"\"timbre\":\"%s\",\"maxDuration\":%d,\"normalize\":\"%s\"}",
|
|
total, json_escape(o.out_dir).c_str(), json_escape(dit_name).c_str(), compat.c_str(), dtype.c_str(),
|
|
timbre.c_str(), o.max_duration, normalize.c_str());
|
|
fprintf(stderr, "[ace-train] preprocess: %d songs -> %s\n", total, o.out_dir.c_str());
|
|
|
|
// ── load models (each stage timed) ──────────────────────────────────
|
|
PreprocessCtx ctx;
|
|
if (!pctx_load(ctx, o, err, jl_stage)) {
|
|
fprintf(stderr, "ace-train preprocess: %s\n", err.c_str());
|
|
jl("{\"type\":\"log\",\"level\":\"error\",\"message\":\"%s\"}", json_escape(err).c_str());
|
|
pctx_free(ctx);
|
|
return 1;
|
|
}
|
|
|
|
// ── driver loop ─────────────────────────────────────────────────────
|
|
PreprocessMeta meta;
|
|
meta.producer = std::string("ace-train ") + ACE_VERSION;
|
|
meta.created_at = pm_iso8601_utc_now();
|
|
meta.manifest = manifest_path;
|
|
meta.dataset_json = manifest_path.empty() ? dataset_path : mf.dataset_json_path;
|
|
meta.audio_dir = mf.source_dir;
|
|
meta.model_variant = dit_name;
|
|
meta.dit_path = o.dit_path;
|
|
meta.vae_path = o.vae_path;
|
|
meta.text_encoder_path = o.text_enc_path;
|
|
meta.compat = compat;
|
|
meta.dtype = (o.dtype == STW_F32) ? "F32" : "BF16";
|
|
meta.timbre = timbre;
|
|
meta.normalize = normalize;
|
|
meta.target_db = o.target_db;
|
|
meta.max_duration = o.max_duration;
|
|
meta.max_caption_tokens = o.max_caption_tokens;
|
|
meta.max_lyric_tokens = o.max_lyric_tokens;
|
|
meta.vae_chunk = o.vae_chunk;
|
|
meta.vae_overlap = o.vae_overlap;
|
|
meta.custom_tag = mf.custom_tag;
|
|
meta.tag_position = mf.tag_position;
|
|
meta.genre_ratio = mf.genre_ratio;
|
|
meta.total = total;
|
|
|
|
double ch_sum[64] = { 0 }, ch_sqsum[64] = { 0 };
|
|
long long n_frames = 0;
|
|
int n_samples = 0;
|
|
int processed = 0, skipped = 0, failed = 0;
|
|
|
|
for (int i = 0; i < total; i++) {
|
|
const PSample & s = mf.samples[(size_t) i];
|
|
const std::string out = o.out_dir + "/" + s.out_stem + ".safetensors";
|
|
|
|
PSampleMeta sm;
|
|
sm.id = s.id;
|
|
sm.rel_path = s.rel_path;
|
|
sm.audio_path = s.audio_path;
|
|
sm.src_bytes = s.src_bytes;
|
|
sm.src_mtime_ms = s.src_mtime_ms;
|
|
|
|
// Resume (P20): skip when the final file already exists — but ONLY when
|
|
// it was produced with the settings this run is using. Skipping on
|
|
// filename existence alone makes preprocess_meta.json (which the server
|
|
// reports to the UI as provenance) claim the CURRENT compat / dtype /
|
|
// timbre / max_duration for files whose own __metadata__ says otherwise;
|
|
// a trainer that trusts the meta then trains on the wrong prompts.
|
|
// A settings change re-encodes; an unchanged re-run is still instant.
|
|
std::string resume_diff;
|
|
if (!o.overwrite && pm_file_exists(out)) {
|
|
PCacheSettings cs;
|
|
if (pm_read_cached_settings(out.c_str(), cs)) {
|
|
auto cmp_s = [&](const char * k, const std::string & had, const std::string & want) {
|
|
if (!had.empty() && had != want) {
|
|
if (!resume_diff.empty()) resume_diff += ", ";
|
|
resume_diff += std::string(k) + " " + had + "\xE2\x86\x92" + want;
|
|
}
|
|
};
|
|
auto cmp_i = [&](const char * k, int had, int want) {
|
|
if (had >= 0 && had != want) {
|
|
char b[96];
|
|
snprintf(b, sizeof(b), "%s %d\xE2\x86\x92%d", k, had, want);
|
|
if (!resume_diff.empty()) resume_diff += ", ";
|
|
resume_diff += b;
|
|
}
|
|
};
|
|
cmp_s("compat", cs.compat, compat);
|
|
cmp_s("dtype", cs.dtype, (o.dtype == STW_F32) ? "F32" : "BF16");
|
|
cmp_s("timbre", cs.timbre, timbre);
|
|
cmp_s("normalize", cs.normalize, normalize);
|
|
cmp_s("model_variant", cs.model_variant, dit_name);
|
|
cmp_i("max_duration", cs.max_duration, o.max_duration);
|
|
cmp_i("max_caption_tokens", cs.max_caption_tokens, o.max_caption_tokens);
|
|
cmp_i("max_lyric_tokens", cs.max_lyric_tokens, o.max_lyric_tokens);
|
|
cmp_i("vae_chunk", cs.vae_chunk, o.vae_chunk);
|
|
cmp_i("vae_overlap", cs.vae_overlap, o.vae_overlap);
|
|
}
|
|
if (!resume_diff.empty()) {
|
|
char wb[640];
|
|
snprintf(wb, sizeof(wb), "%s: cache was built with different settings (%s) - re-encoding",
|
|
s.rel_path.c_str(), resume_diff.c_str());
|
|
fprintf(stderr, "[ace-train] %s\n", wb);
|
|
jl("{\"type\":\"log\",\"level\":\"warn\",\"message\":\"%s\"}", json_escape(wb).c_str());
|
|
}
|
|
}
|
|
|
|
if (!o.overwrite && resume_diff.empty() && pm_file_exists(out)) {
|
|
std::vector<float> cached;
|
|
int T = 0;
|
|
if (pm_read_cached_latents(out.c_str(), cached, &T, &sm) && T > 0) {
|
|
for (int t = 0; t < T; t++) {
|
|
const float * row = cached.data() + (size_t) t * 64;
|
|
for (int ch = 0; ch < 64; ch++) {
|
|
ch_sum[ch] += row[ch];
|
|
ch_sqsum[ch] += (double) row[ch] * (double) row[ch];
|
|
}
|
|
}
|
|
n_frames += T;
|
|
n_samples++;
|
|
}
|
|
long long bytes = 0;
|
|
pm_stat_file(out, &bytes, NULL);
|
|
sm.file = pm_basename(out);
|
|
sm.bytes = bytes;
|
|
sm.ok = true;
|
|
sm.skipped = true;
|
|
meta.samples.push_back(sm);
|
|
skipped++;
|
|
jl("{\"type\":\"song_skip\",\"index\":%d,\"id\":\"%s\",\"file\":\"%s\",\"reason\":\"exists\"}", i,
|
|
json_escape(s.id).c_str(), json_escape(s.rel_path).c_str());
|
|
jl("{\"type\":\"progress\",\"completed\":%d,\"total\":%d,\"processed\":%d,\"skipped\":%d,\"failed\":%d}",
|
|
processed + skipped + failed, total, processed, skipped, failed);
|
|
continue;
|
|
}
|
|
|
|
// A song that cannot be attempted is a SKIP, not a failure (§2.2). Left
|
|
// as song_fail these count toward `failed`, and an all-FLAC dataset run
|
|
// without --ffmpeg exits 1 — which the server surfaces as a hard job
|
|
// failure with no hint that the cause is per-song and per-format.
|
|
if (const char * skip_reason = pp_skip_reason(o, s)) {
|
|
sm.file = "";
|
|
sm.ok = false;
|
|
sm.skipped = true;
|
|
sm.error = (strcmp(skip_reason, "missing-file") == 0)
|
|
? ("audio file not found: " + s.audio_path)
|
|
: ("unsupported audio format " + pm_ext_of(s.audio_path) + "; ffmpeg not configured");
|
|
meta.samples.push_back(sm);
|
|
skipped++;
|
|
fprintf(stderr, "[ace-train] SKIP %s: %s\n", s.rel_path.c_str(), sm.error.c_str());
|
|
jl("{\"type\":\"song_skip\",\"index\":%d,\"id\":\"%s\",\"file\":\"%s\",\"reason\":\"%s\"}", i,
|
|
json_escape(s.id).c_str(), json_escape(s.rel_path).c_str(), skip_reason);
|
|
jl("{\"type\":\"log\",\"level\":\"warn\",\"message\":\"%s\"}", json_escape(sm.error).c_str());
|
|
jl("{\"type\":\"progress\",\"completed\":%d,\"total\":%d,\"processed\":%d,\"skipped\":%d,\"failed\":%d}",
|
|
processed + skipped + failed, total, processed, skipped, failed);
|
|
continue;
|
|
}
|
|
|
|
jl("{\"type\":\"song_start\",\"index\":%d,\"id\":\"%s\",\"file\":\"%s\"}", i, json_escape(s.id).c_str(),
|
|
json_escape(s.rel_path).c_str());
|
|
fprintf(stderr, "[ace-train] [%d/%d] %s\n", i + 1, total, s.rel_path.c_str());
|
|
|
|
const int64_t t0 = ggml_time_ms();
|
|
SongResult r = preprocess_song(ctx, o, mf, s, out, jl_warn);
|
|
const int64_t ms = ggml_time_ms() - t0;
|
|
|
|
if (r.ok) {
|
|
for (int t = 0; t < r.t_latent; t++) {
|
|
const float * row = r.latents.data() + (size_t) t * 64;
|
|
for (int ch = 0; ch < 64; ch++) {
|
|
ch_sum[ch] += row[ch];
|
|
ch_sqsum[ch] += (double) row[ch] * (double) row[ch];
|
|
}
|
|
}
|
|
n_frames += r.t_latent;
|
|
n_samples++;
|
|
processed++;
|
|
|
|
sm.file = pm_basename(out);
|
|
sm.bytes = r.bytes;
|
|
sm.t_latent = r.t_latent;
|
|
sm.s_total = r.s_total;
|
|
sm.s_text = r.s_text;
|
|
sm.s_lyric = r.s_lyric;
|
|
sm.s_total_genre = r.s_total_genre;
|
|
sm.has_genre = r.has_genre;
|
|
sm.ok = true;
|
|
meta.samples.push_back(sm);
|
|
|
|
jl("{\"type\":\"song_done\",\"index\":%d,\"id\":\"%s\",\"file\":\"%s\",\"out\":\"%s\",\"tLatent\":%d,"
|
|
"\"sTotal\":%d,\"sText\":%d,\"sLyric\":%d,\"sTotalGenre\":%d,\"hasGenre\":%s,\"bytes\":%lld,"
|
|
"\"ms\":%lld}",
|
|
i, json_escape(s.id).c_str(), json_escape(s.rel_path).c_str(), json_escape(sm.file).c_str(), r.t_latent,
|
|
r.s_total, r.s_text, r.s_lyric, r.s_total_genre, r.has_genre ? "true" : "false", (long long) r.bytes,
|
|
(long long) ms);
|
|
} else {
|
|
failed++;
|
|
sm.ok = false;
|
|
sm.error = r.error;
|
|
meta.samples.push_back(sm);
|
|
fprintf(stderr, "[ace-train] FAILED %s: %s\n", s.rel_path.c_str(), r.error.c_str());
|
|
jl("{\"type\":\"song_fail\",\"index\":%d,\"id\":\"%s\",\"file\":\"%s\",\"error\":\"%s\"}", i,
|
|
json_escape(s.id).c_str(), json_escape(s.rel_path).c_str(), json_escape(r.error).c_str());
|
|
}
|
|
|
|
jl("{\"type\":\"progress\",\"completed\":%d,\"total\":%d,\"processed\":%d,\"skipped\":%d,\"failed\":%d}",
|
|
processed + skipped + failed, total, processed, skipped, failed);
|
|
}
|
|
|
|
// Free the GPU as early as possible — before the stats/meta writes.
|
|
pctx_free(ctx);
|
|
|
|
meta.processed = processed;
|
|
meta.skipped = skipped;
|
|
meta.failed = failed;
|
|
|
|
if (n_frames > 0) {
|
|
const std::string stats_path = o.out_dir + "/channel_stats.json";
|
|
if (pm_write_channel_stats(stats_path.c_str(), ch_sum, ch_sqsum, n_frames, n_samples)) {
|
|
jl("{\"type\":\"stats\",\"nSamples\":%d,\"nFrames\":%lld}", n_samples, (long long) n_frames);
|
|
} else {
|
|
fprintf(stderr, "[ace-train] WARNING: failed to write %s\n", stats_path.c_str());
|
|
}
|
|
}
|
|
|
|
const std::string meta_path = o.out_dir + "/preprocess_meta.json";
|
|
if (!pm_write_meta_json(meta_path.c_str(), meta)) {
|
|
fprintf(stderr, "[ace-train] FATAL: failed to write %s\n", meta_path.c_str());
|
|
return 1;
|
|
}
|
|
|
|
const long long run_ms = (long long) (ggml_time_ms() - t_run0);
|
|
jl("{\"type\":\"done\",\"processed\":%d,\"skipped\":%d,\"failed\":%d,\"out\":\"%s\",\"ms\":%lld}", processed,
|
|
skipped, failed, json_escape(o.out_dir).c_str(), run_ms);
|
|
fprintf(stderr, "[ace-train] done: %d processed, %d skipped, %d failed (%.1f s)\n", processed, skipped, failed,
|
|
(double) run_ms / 1000.0);
|
|
|
|
if (!g_jsonl) {
|
|
printf("preprocess: %d processed, %d skipped, %d failed -> %s\n", processed, skipped, failed, o.out_dir.c_str());
|
|
fflush(stdout);
|
|
}
|
|
|
|
const int attempted = processed + failed;
|
|
return (processed == 0 && attempted > 0) ? 1 : 0;
|
|
}
|
|
|
|
// ─── --bwd: MUL_MAT activation-gradient formulation ─────────────────────────
|
|
//
|
|
// engine/patches/mm-backward.patch teaches ggml_compute_backward a second way to
|
|
// emit the gradient wrt src1 of a MUL_MAT:
|
|
// mm : mul_mat(cont(transpose(src0)), grad)
|
|
// outprod : out_prod(src0, transpose(grad)) <- upstream ggml
|
|
// The two are shape- and maths-identical, but out_prod is F32-only on CUDA
|
|
// (cublasSgemm), which forces the frozen weight to F32 and drags the FORWARD
|
|
// mul_mat onto TF32 as well. The mul_mat arm is dtype-agnostic, so a BF16 weight
|
|
// rides real BF16 tensor cores in both directions. Measured on an RTX 5090:
|
|
// ~1.7-1.8x per layer per step (engine/src/train/spike-gemmbench.h).
|
|
//
|
|
// The patch reads GGML_BACKWARD_MM ONCE, at the first backward graph build, so
|
|
// this MUST run before any model load or graph construction — hence it lives in
|
|
// the argv handlers rather than inside the trainers, and covers --self-test too.
|
|
//
|
|
// `outprod` deliberately does NOT clear the variable. The self-test battery is
|
|
// A/B'd by exporting GGML_BACKWARD_MM in the environment, and clearing it here
|
|
// would make the default `--bwd outprod` silently defeat that.
|
|
static bool bwd_valid(const std::string & v) {
|
|
return v == "outprod" || v == "mm";
|
|
}
|
|
|
|
static void bwd_apply(const std::string & v) {
|
|
if (v != "mm") {
|
|
return;
|
|
}
|
|
#ifdef _WIN32
|
|
_putenv("GGML_BACKWARD_MM=1");
|
|
#else
|
|
setenv("GGML_BACKWARD_MM", "1", 1);
|
|
#endif
|
|
fprintf(stderr, "[ace-train] --bwd mm: GGML_BACKWARD_MM=1 (mul_mat activation backward)\n");
|
|
}
|
|
|
|
// ─── train-lm ───────────────────────────────────────────────────────────────
|
|
//
|
|
// docs/plans/2026-07-27-lm-trainer-implementation.md §2.1, §3.1
|
|
|
|
static int cmd_train_lm(int argc, char ** argv) {
|
|
LmTrainArgs a;
|
|
std::string stages_csv, dit_arg, lm_arg, lm_size_arg;
|
|
|
|
for (int i = 1; i < argc; i++) {
|
|
if (!strcmp(argv[i], "--stages") && i + 1 < argc) stages_csv = argv[++i];
|
|
else if (!strcmp(argv[i], "--tensors") && i + 1 < argc) a.tensors_dir = argv[++i];
|
|
else if (!strcmp(argv[i], "--codes") && i + 1 < argc) a.codes_path = argv[++i];
|
|
else if (!strcmp(argv[i], "--out") && i + 1 < argc) a.out_dir = argv[++i];
|
|
else if (!strcmp(argv[i], "--models") && i + 1 < argc) a.models_dir = argv[++i];
|
|
else if (!strcmp(argv[i], "--dit") && i + 1 < argc) dit_arg = argv[++i];
|
|
else if (!strcmp(argv[i], "--lm") && i + 1 < argc) lm_arg = argv[++i];
|
|
else if (!strcmp(argv[i], "--lm-size") && i + 1 < argc) lm_size_arg = argv[++i];
|
|
else if (!strcmp(argv[i], "--rank") && i + 1 < argc) a.rank = atoi(argv[++i]);
|
|
else if (!strcmp(argv[i], "--alpha") && i + 1 < argc) a.alpha = atoi(argv[++i]);
|
|
else if (!strcmp(argv[i], "--lr") && i + 1 < argc) a.lr = (float) atof(argv[++i]);
|
|
else if (!strcmp(argv[i], "--epochs") && i + 1 < argc) a.epochs = atoi(argv[++i]);
|
|
else if (!strcmp(argv[i], "--grad-accum") && i + 1 < argc) a.grad_accum = atoi(argv[++i]);
|
|
else if (!strcmp(argv[i], "--warmup-ratio") && i + 1 < argc) a.warmup_ratio = (float) atof(argv[++i]);
|
|
else if (!strcmp(argv[i], "--grad-clip") && i + 1 < argc) a.grad_clip = (float) atof(argv[++i]);
|
|
else if (!strcmp(argv[i], "--weight-decay") && i + 1 < argc) a.weight_decay = (float) atof(argv[++i]);
|
|
else if (!strcmp(argv[i], "--seed") && i + 1 < argc) a.seed = atoi(argv[++i]);
|
|
else if (!strcmp(argv[i], "--target-loss") && i + 1 < argc) a.target_loss = (float) atof(argv[++i]);
|
|
else if (!strcmp(argv[i], "--order") && i + 1 < argc) a.order = argv[++i];
|
|
else if (!strcmp(argv[i], "--max-len") && i + 1 < argc) a.max_len = atoi(argv[++i]);
|
|
else if (!strcmp(argv[i], "--vram-reserve-mb") && i + 1 < argc) a.vram_reserve_mb = atoi(argv[++i]);
|
|
else if (!strcmp(argv[i], "--low-vram") && i + 1 < argc) a.low_vram = argv[++i];
|
|
else if (!strcmp(argv[i], "--attn-head-block") && i + 1 < argc) a.attn_head_block = atoi(argv[++i]);
|
|
else if (!strcmp(argv[i], "--lm-chunk") && i + 1 < argc) a.chunk = atoi(argv[++i]);
|
|
else if (!strcmp(argv[i], "--weights") && i + 1 < argc) a.weights = argv[++i];
|
|
else if (!strcmp(argv[i], "--optimizer") && i + 1 < argc) a.optimizer = argv[++i];
|
|
else if (!strcmp(argv[i], "--muon-lr-scale") && i + 1 < argc) a.muon_lr_scale = (float) atof(argv[++i]);
|
|
else if (!strcmp(argv[i], "--muon-momentum") && i + 1 < argc) a.muon_momentum = (float) atof(argv[++i]);
|
|
else if (!strcmp(argv[i], "--muon-ns-steps") && i + 1 < argc) a.muon_ns_steps = atoi(argv[++i]);
|
|
else if (!strcmp(argv[i], "--muon-min-dim") && i + 1 < argc) a.muon_min_dim = atoi(argv[++i]);
|
|
else if (!strcmp(argv[i], "--muon-bucket") && i + 1 < argc) a.muon_bucket = atoi(argv[++i]);
|
|
else if (!strcmp(argv[i], "--no-muon-nesterov")) a.muon_nesterov = false;
|
|
else if (!strcmp(argv[i], "--adapter-type") && i + 1 < argc) a.adapter_type = argv[++i];
|
|
else if (!strcmp(argv[i], "--lokr-dim") && i + 1 < argc) a.lokr_dim = atoi(argv[++i]);
|
|
else if (!strcmp(argv[i], "--lokr-alpha") && i + 1 < argc) a.lokr_alpha = (float) atof(argv[++i]);
|
|
else if (!strcmp(argv[i], "--lokr-factor") && i + 1 < argc) a.lokr_factor = atoi(argv[++i]);
|
|
else if (!strcmp(argv[i], "--no-lokr-decompose-both")) a.lokr_decompose_both = false;
|
|
else if (!strcmp(argv[i], "--bwd") && i + 1 < argc) a.bwd = argv[++i];
|
|
else if (!strcmp(argv[i], "--batch") && i + 1 < argc) a.batch = argv[++i];
|
|
else if (!strcmp(argv[i], "--trigger") && i + 1 < argc) a.trigger = argv[++i];
|
|
else if (!strcmp(argv[i], "--trigger-position") && i + 1 < argc) a.trigger_position = argv[++i];
|
|
else if (!strcmp(argv[i], "--milestone-step") && i + 1 < argc) a.milestone_step = (float) atof(argv[++i]);
|
|
else if (!strcmp(argv[i], "--milestone-keep") && i + 1 < argc) a.milestone_keep = atoi(argv[++i]);
|
|
else if (!strcmp(argv[i], "--limit") && i + 1 < argc) a.limit = atoi(argv[++i]);
|
|
else if (!strcmp(argv[i], "--loss-on-cot")) a.loss_on_cot = true;
|
|
else if (!strcmp(argv[i], "--no-loss-on-cot")) a.loss_on_cot = false;
|
|
else if (!strcmp(argv[i], "--no-milestones")) a.milestone_step = 0.0f;
|
|
else if (!strcmp(argv[i], "--overwrite")) a.overwrite = true;
|
|
else if (!strcmp(argv[i], "--self-test")) a.self_test = true;
|
|
else if (!strcmp(argv[i], "--jsonl")) g_jsonl = true;
|
|
else if (!strcmp(argv[i], "-h") || !strcmp(argv[i], "--help")) { print_usage(); return 0; }
|
|
else {
|
|
fprintf(stderr, "ace-train train-lm: unknown option '%s'\n", argv[i]);
|
|
print_usage();
|
|
return 2;
|
|
}
|
|
}
|
|
|
|
// ── stages ──────────────────────────────────────────────────────────
|
|
if (!stages_csv.empty()) {
|
|
const bool want_ex = stages_csv.find("extract") != std::string::npos;
|
|
const bool want_tr = stages_csv.find("train") != std::string::npos;
|
|
const bool want_xp = stages_csv.find("export") != std::string::npos;
|
|
a.stages.clear();
|
|
if (want_ex) a.stages.push_back("extract");
|
|
if (want_tr) a.stages.push_back("train");
|
|
if (want_xp) a.stages.push_back("export");
|
|
if (a.stages.empty()) {
|
|
fprintf(stderr, "ace-train train-lm: --stages must name at least one of extract,train,export\n");
|
|
return 2;
|
|
}
|
|
}
|
|
|
|
// ── numeric sanity (the server clamps these too; be defensive) ───────
|
|
if (a.rank < 1 || a.rank > 256 || a.alpha < 1 || a.alpha > 1024 || a.epochs < 1 || a.epochs > 200 ||
|
|
a.grad_accum < 1 || a.grad_accum > 64 || a.lr <= 0.0f || a.lr > 1.0f || a.grad_clip < 0.0f ||
|
|
a.warmup_ratio < 0.0f || a.warmup_ratio > 0.5f || a.weight_decay < 0.0f || a.weight_decay > 1.0f ||
|
|
a.target_loss < 0.0f || a.target_loss > 20.0f || a.milestone_keep < 0 || a.milestone_keep > 64 ||
|
|
a.vram_reserve_mb < 0 || (a.max_len != 0 && (a.max_len < 512 || a.max_len > 16384))) {
|
|
fprintf(stderr, "ace-train train-lm: numeric option out of range\n");
|
|
return 2;
|
|
}
|
|
if (a.order != "shuffle" && a.order != "fixed") {
|
|
fprintf(stderr, "ace-train train-lm: --order must be shuffle|fixed\n");
|
|
return 2;
|
|
}
|
|
// "replace" never applies the tag during preprocessing (preprocess-run.h:203),
|
|
// so an adapter trained from it carried no trigger and must not claim one.
|
|
if (!a.trigger_position.empty() && a.trigger_position != "prepend" && a.trigger_position != "append") {
|
|
fprintf(stderr, "ace-train train-lm: --trigger-position must be prepend|append\n");
|
|
return 2;
|
|
}
|
|
if (a.trigger.size() > 128) {
|
|
fprintf(stderr, "ace-train train-lm: --trigger must be 128 characters or fewer\n");
|
|
return 2;
|
|
}
|
|
if (a.low_vram != "auto" && a.low_vram != "on" && a.low_vram != "off") {
|
|
fprintf(stderr, "ace-train train-lm: --low-vram must be auto|on|off\n");
|
|
return 2;
|
|
}
|
|
if (a.attn_head_block < -1 || a.attn_head_block > 128 || a.chunk < 16 || a.chunk > 1024) {
|
|
fprintf(stderr, "ace-train train-lm: --attn-head-block must be -1..128 and --lm-chunk 16..1024\n");
|
|
return 2;
|
|
}
|
|
|
|
// ── speed levers (2026-07-28 plan §2.1) ─────────────────────────────
|
|
//
|
|
// Every one of these is raised BEFORE any model load, so the server never
|
|
// pays an engine stop/restart for a bad flag.
|
|
if (a.weights != "f32-window" && a.weights != "bf16") {
|
|
fprintf(stderr, "ace-train train-lm: --weights must be f32-window|bf16\n");
|
|
return 2;
|
|
}
|
|
if (!bwd_valid(a.bwd)) {
|
|
fprintf(stderr, "ace-train train-lm: --bwd must be outprod|mm\n");
|
|
return 2;
|
|
}
|
|
// COLLISION, refused rather than coerced.
|
|
//
|
|
// `--weights bf16` (lm-bf16.h, Lever A) reaches the SAME mul_mat backward by
|
|
// a different route: it lets ggml build the out_prod form and then rewrites
|
|
// those nodes in place, asserting EXACTLY 7 rewrites / 0 skipped / 0 residual
|
|
// per segment (its S18 tripwire). `--bwd mm` makes ggml emit mul_mat directly,
|
|
// so the surgery finds nothing to rewrite and the tripwire GGML_ABORTs — which
|
|
// is the tripwire doing its job, but it would abort mid-run, after the model
|
|
// load, rather than here.
|
|
//
|
|
// There is no version of this pair worth building: on the f32-window path
|
|
// `--bwd mm` has nothing to gain either, because the weight it transposes is
|
|
// the F32 window, so the GEMM stays F32/TF32 and the extra `cont` is pure
|
|
// cost. The LM already solved this problem its own way; `--bwd mm` is a win
|
|
// for train-dit, whose bf16 mirror has no such surgery.
|
|
if (a.weights == "bf16" && a.bwd == "mm") {
|
|
fprintf(stderr,
|
|
"ace-train train-lm: --weights bf16 and --bwd mm are two routes to the SAME\n"
|
|
" mul_mat backward and cannot be combined. --weights bf16 rewrites ggml's out_prod\n"
|
|
" nodes in place and asserts it found exactly 7 per segment (lm-bf16.h S18); --bwd mm\n"
|
|
" makes ggml emit mul_mat directly, so that surgery finds nothing and aborts the run.\n"
|
|
" Use --weights bf16 (it already gives you the BF16 tensor-core backward) or pair\n"
|
|
" --bwd mm with --weights f32-window.\n");
|
|
return 2;
|
|
}
|
|
int batch_n = 1;
|
|
if (a.batch != "auto") {
|
|
char * end = nullptr;
|
|
const long bv = strtol(a.batch.c_str(), &end, 10);
|
|
if (!end || *end != '\0' || bv < 1 || bv > 8) {
|
|
fprintf(stderr, "ace-train train-lm: --batch must be 1..8 or auto\n");
|
|
return 2;
|
|
}
|
|
batch_n = (int) bv;
|
|
}
|
|
// bf16 has no F32 mirror to fall back on, so it is only meaningful under
|
|
// checkpointing. An explicit `--low-vram off` is a direct contradiction,
|
|
// not something to silently override.
|
|
if (a.weights == "bf16" && a.low_vram == "off") {
|
|
fprintf(stderr, "ace-train train-lm: --weights bf16 requires low-VRAM mode "
|
|
"(the naive path mirrors every weight to F32) — drop --low-vram off\n");
|
|
return 2;
|
|
}
|
|
// LEVER B IS NOT BUILT — see the §6.1 Phase-0 measurement, and read the
|
|
// numbers rather than the headline, because they are more equivocal than a
|
|
// one-word verdict suggests. Batching cannot make the GEMMs cheaper (S7);
|
|
// its whole win is spreading per-micro-step HOST cost over B samples, so the
|
|
// design made it conditional on that cost reaching 10 %. Measured
|
|
// (HOTSTEP_LM_PHASE_TIMER=1, 20 micro-steps, S=2427):
|
|
//
|
|
// 4B low-VRAM: build 2.35 % + reset 0.01 % + plan 0.20 % + submit 6.75 %
|
|
// = 9.31 % — under the bar even counting all of submit.
|
|
// 0.6B low-VRAM: 7.43 + 0.03 + 0.35 + 12.43 = 20.24 %, BUT `submit` is an
|
|
// upper bound (async launches absorb GPU wait when the queue
|
|
// backs up), so the true figure is bracketed [7.81 %,
|
|
// 20.24 %] — the bar sits INSIDE the bracket. Inconclusive,
|
|
// not a miss.
|
|
//
|
|
// So: definitively closed at 4B (the case that costs real wall clock), and
|
|
// unresolved at 0.6B low-VRAM — which is the narrowest slice there is, since
|
|
// 0.6B runs naive by default and is FASTER that way (257 vs 368 ms). Not
|
|
// worth ~600 lines through the hottest code in the trainer on that basis.
|
|
// Accepting the flag and silently packing batches of 1 is exactly the
|
|
// failure mode S18 forbids, so this refuses loudly instead.
|
|
//
|
|
// THIS CHECK RUNS BEFORE THE --low-vram INTERACTION BELOW, deliberately.
|
|
// `--batch 2 --low-vram off` used to answer "--batch >1 requires low-VRAM
|
|
// mode", which sends the user off to fix something that still would not
|
|
// work. The unconditional refusal is the true reason, so it speaks first.
|
|
if (a.batch == "auto" || batch_n > 1) {
|
|
fprintf(stderr,
|
|
"ace-train train-lm: --batch %s is not supported in this build.\n"
|
|
" Micro-batching was never written. It cannot make the maths cheaper — it only\n"
|
|
" spreads per-micro-step host overhead over B samples, and that overhead measures\n"
|
|
" 9.3%% at 4B (under the 10%% bar the design required) and at most 20%% at 0.6B in\n"
|
|
" low-VRAM mode, which is a mode 0.6B does not need. Use --grad-accum instead — it\n"
|
|
" changes the effective batch size with none of the padding waste.\n",
|
|
a.batch.c_str());
|
|
return 2;
|
|
}
|
|
// Kept for the day Lever B is built: at that point this is the rule, and
|
|
// until then it is unreachable-by-construction rather than wrong.
|
|
if ((a.batch == "auto" || batch_n > 1) && a.low_vram == "off") {
|
|
fprintf(stderr, "ace-train train-lm: --batch >1 requires low-VRAM mode — drop --low-vram off\n");
|
|
return 2;
|
|
}
|
|
|
|
// ── model resolution (same rule as preprocess) ──────────────────────
|
|
ModelRegistry reg;
|
|
if (!a.models_dir.empty()) {
|
|
registry_scan(®, a.models_dir.c_str());
|
|
}
|
|
if (!lm_arg.empty()) {
|
|
std::string name;
|
|
if (!resolve_model(lm_arg, reg.lm, false, a.lm_path, name)) {
|
|
fprintf(stderr, "ace-train train-lm: cannot resolve --lm '%s' in %s\n", lm_arg.c_str(),
|
|
a.models_dir.c_str());
|
|
return 2;
|
|
}
|
|
a.lm_name = name;
|
|
}
|
|
if (!dit_arg.empty()) {
|
|
std::string name;
|
|
if (!resolve_model(dit_arg, reg.dit, false, a.dit_path, name)) {
|
|
fprintf(stderr, "ace-train train-lm: cannot resolve --dit '%s' in %s\n", dit_arg.c_str(),
|
|
a.models_dir.c_str());
|
|
return 2;
|
|
}
|
|
}
|
|
|
|
// ── required arguments per stage ────────────────────────────────────
|
|
if (a.self_test) {
|
|
if (a.lm_path.empty()) {
|
|
fprintf(stderr, "ace-train train-lm --self-test: --lm is required\n");
|
|
return 2;
|
|
}
|
|
} else {
|
|
if (lm_has_stage(a, "extract") && a.tensors_dir.empty()) {
|
|
fprintf(stderr, "ace-train train-lm: --tensors is required for the extract stage\n");
|
|
return 2;
|
|
}
|
|
if (lm_has_stage(a, "train") && a.lm_path.empty()) {
|
|
fprintf(stderr, "ace-train train-lm: --lm is required for the train stage\n");
|
|
return 2;
|
|
}
|
|
if ((lm_has_stage(a, "train") || lm_has_stage(a, "export")) && a.out_dir.empty()) {
|
|
fprintf(stderr, "ace-train train-lm: --out is required for the train/export stages\n");
|
|
return 2;
|
|
}
|
|
}
|
|
if (a.codes_path.empty()) {
|
|
if (a.tensors_dir.empty()) {
|
|
if (!a.self_test) {
|
|
fprintf(stderr, "ace-train train-lm: --codes or --tensors is required\n");
|
|
return 2;
|
|
}
|
|
} else {
|
|
a.codes_path = lm_join(a.tensors_dir, "lm_codes.jsonl");
|
|
}
|
|
}
|
|
|
|
// ── lm_size label ───────────────────────────────────────────────────
|
|
if (!lm_size_arg.empty()) {
|
|
a.lm_size = lm_size_arg;
|
|
} else if (!a.lm_path.empty()) {
|
|
const Qwen3LMConfig c = qw3lm_load_config(a.lm_path.c_str(), !ends_with_gguf(a.lm_path.c_str()));
|
|
a.lm_size = lm_size_label_from_config(c);
|
|
}
|
|
|
|
// MUST precede every model load and graph build — the ggml patch latches
|
|
// GGML_BACKWARD_MM on the first backward it emits.
|
|
bwd_apply(a.bwd);
|
|
|
|
// MANDATORY: ggml_time_ms() divides by an uninitialised frequency otherwise.
|
|
ggml_time_init();
|
|
return lm_train_main(a);
|
|
}
|
|
|
|
// ─── train-dit (plan §2.1) ──────────────────────────────────────────────────
|
|
|
|
static int cmd_train_dit(int argc, char ** argv) {
|
|
DitTrainArgs a;
|
|
std::string stages_csv, dit_arg;
|
|
bool safety_user = false;
|
|
|
|
for (int i = 1; i < argc; i++) {
|
|
if (!strcmp(argv[i], "--stages") && i + 1 < argc) stages_csv = argv[++i];
|
|
else if (!strcmp(argv[i], "--tensors") && i + 1 < argc) a.tensors_dir = argv[++i];
|
|
else if (!strcmp(argv[i], "--out") && i + 1 < argc) a.out_dir = argv[++i];
|
|
else if (!strcmp(argv[i], "--models") && i + 1 < argc) a.models_dir = argv[++i];
|
|
else if (!strcmp(argv[i], "--dit") && i + 1 < argc) dit_arg = argv[++i];
|
|
else if (!strcmp(argv[i], "--adapter-type") && i + 1 < argc) a.adapter_type = argv[++i];
|
|
else if (!strcmp(argv[i], "--rank") && i + 1 < argc) a.rank = atoi(argv[++i]);
|
|
else if (!strcmp(argv[i], "--alpha") && i + 1 < argc) a.alpha = atoi(argv[++i]);
|
|
else if (!strcmp(argv[i], "--lokr-dim") && i + 1 < argc) a.lokr_dim = atoi(argv[++i]);
|
|
else if (!strcmp(argv[i], "--lokr-alpha") && i + 1 < argc) a.lokr_alpha = (float) atof(argv[++i]);
|
|
else if (!strcmp(argv[i], "--lokr-factor") && i + 1 < argc) a.lokr_factor = atoi(argv[++i]);
|
|
else if (!strcmp(argv[i], "--lokr-decompose-both")) a.lokr_decompose_both = true;
|
|
else if (!strcmp(argv[i], "--no-lokr-decompose-both")) a.lokr_decompose_both = false;
|
|
else if (!strcmp(argv[i], "--layers") && i + 1 < argc) a.layers = atoi(argv[++i]);
|
|
else if (!strcmp(argv[i], "--target-mlp")) a.target_mlp = true;
|
|
else if (!strcmp(argv[i], "--no-target-mlp")) a.target_mlp = false;
|
|
else if (!strcmp(argv[i], "--loss-weighting") && i + 1 < argc) a.loss_weighting = argv[++i];
|
|
else if (!strcmp(argv[i], "--snr-gamma") && i + 1 < argc) a.snr_gamma = (float) atof(argv[++i]);
|
|
else if (!strcmp(argv[i], "--t-bias") && i + 1 < argc) a.t_bias = (float) atof(argv[++i]);
|
|
else if (!strcmp(argv[i], "--channel-balance")) a.channel_balance = true;
|
|
else if (!strcmp(argv[i], "--no-channel-balance")) a.channel_balance = false;
|
|
else if (!strcmp(argv[i], "--timestep-mu") && i + 1 < argc) a.timestep_mu = (float) atof(argv[++i]);
|
|
else if (!strcmp(argv[i], "--timestep-sigma") && i + 1 < argc) a.timestep_sigma = (float) atof(argv[++i]);
|
|
else if (!strcmp(argv[i], "--t-min") && i + 1 < argc) a.t_min = (float) atof(argv[++i]);
|
|
else if (!strcmp(argv[i], "--t-max") && i + 1 < argc) a.t_max = (float) atof(argv[++i]);
|
|
else if (!strcmp(argv[i], "--cfg-ratio") && i + 1 < argc) a.cfg_ratio = (float) atof(argv[++i]);
|
|
else if (!strcmp(argv[i], "--genre-ratio") && i + 1 < argc) a.genre_ratio = atoi(argv[++i]);
|
|
else if (!strcmp(argv[i], "--lr") && i + 1 < argc) a.lr = (float) atof(argv[++i]);
|
|
else if (!strcmp(argv[i], "--epochs") && i + 1 < argc) a.epochs = atoi(argv[++i]);
|
|
else if (!strcmp(argv[i], "--grad-accum") && i + 1 < argc) a.grad_accum = atoi(argv[++i]);
|
|
else if (!strcmp(argv[i], "--warmup-ratio") && i + 1 < argc) a.warmup_ratio = (float) atof(argv[++i]);
|
|
else if (!strcmp(argv[i], "--grad-clip") && i + 1 < argc) a.grad_clip = (float) atof(argv[++i]);
|
|
else if (!strcmp(argv[i], "--weight-decay") && i + 1 < argc) a.weight_decay = (float) atof(argv[++i]);
|
|
else if (!strcmp(argv[i], "--seed") && i + 1 < argc) a.seed = atoi(argv[++i]);
|
|
else if (!strcmp(argv[i], "--target-loss") && i + 1 < argc) a.target_loss = (float) atof(argv[++i]);
|
|
else if (!strcmp(argv[i], "--order") && i + 1 < argc) a.order = argv[++i];
|
|
else if (!strcmp(argv[i], "--batch") && i + 1 < argc) a.batch = atoi(argv[++i]);
|
|
else if (!strcmp(argv[i], "--ckpt") && i + 1 < argc) a.ckpt = atoi(argv[++i]);
|
|
else if (!strcmp(argv[i], "--crop") && i + 1 < argc) a.crop = atoi(argv[++i]);
|
|
else if (!strcmp(argv[i], "--crop-min") && i + 1 < argc) a.crop_min = atoi(argv[++i]);
|
|
else if (!strcmp(argv[i], "--crop-max") && i + 1 < argc) a.crop_max = atoi(argv[++i]);
|
|
else if (!strcmp(argv[i], "--vram-reserve-mb") && i + 1 < argc) a.vram_reserve_mb = atoi(argv[++i]);
|
|
else if (!strcmp(argv[i], "--vram-safety") && i + 1 < argc) { a.vram_safety = (float) atof(argv[++i]); safety_user = true; }
|
|
else if (!strcmp(argv[i], "--mirror") && i + 1 < argc) a.mirror = argv[++i];
|
|
else if (!strcmp(argv[i], "--bwd") && i + 1 < argc) a.bwd = argv[++i];
|
|
else if (!strcmp(argv[i], "--profile-step") && i + 1 < argc) a.profile_step = atoi(argv[++i]);
|
|
else if (!strcmp(argv[i], "--profile-ops")) a.profile_ops = true;
|
|
else if (!strcmp(argv[i], "--optimizer") && i + 1 < argc) a.optimizer = argv[++i];
|
|
else if (!strcmp(argv[i], "--muon-lr-scale") && i + 1 < argc) a.muon_lr_scale = (float) atof(argv[++i]);
|
|
else if (!strcmp(argv[i], "--muon-momentum") && i + 1 < argc) a.muon_momentum = (float) atof(argv[++i]);
|
|
else if (!strcmp(argv[i], "--muon-ns-steps") && i + 1 < argc) a.muon_ns_steps = atoi(argv[++i]);
|
|
else if (!strcmp(argv[i], "--muon-min-dim") && i + 1 < argc) a.muon_min_dim = atoi(argv[++i]);
|
|
else if (!strcmp(argv[i], "--muon-bucket") && i + 1 < argc) a.muon_bucket = atoi(argv[++i]);
|
|
else if (!strcmp(argv[i], "--no-muon-nesterov")) a.muon_nesterov = false;
|
|
else if (!strcmp(argv[i], "--trigger") && i + 1 < argc) a.trigger = argv[++i];
|
|
else if (!strcmp(argv[i], "--trigger-position") && i + 1 < argc) a.trigger_position = argv[++i];
|
|
else if (!strcmp(argv[i], "--milestone-step") && i + 1 < argc) a.milestone_step = (float) atof(argv[++i]);
|
|
else if (!strcmp(argv[i], "--milestone-keep") && i + 1 < argc) a.milestone_keep = atoi(argv[++i]);
|
|
else if (!strcmp(argv[i], "--no-milestones")) a.milestone_step = 0.0f;
|
|
else if (!strcmp(argv[i], "--limit") && i + 1 < argc) a.limit = atoi(argv[++i]);
|
|
else if (!strcmp(argv[i], "--overwrite")) a.overwrite = true;
|
|
else if (!strcmp(argv[i], "--self-test")) a.self_test = true;
|
|
else if (!strcmp(argv[i], "--jsonl")) g_jsonl = true;
|
|
else if (!strcmp(argv[i], "-h") || !strcmp(argv[i], "--help")) { print_usage(); return 0; }
|
|
else {
|
|
fprintf(stderr, "ace-train train-dit: unknown option '%s'\n", argv[i]);
|
|
print_usage();
|
|
return 2;
|
|
}
|
|
}
|
|
|
|
// ── stages ──────────────────────────────────────────────────────────
|
|
if (!stages_csv.empty()) {
|
|
const bool want_tr = stages_csv.find("train") != std::string::npos;
|
|
const bool want_xp = stages_csv.find("export") != std::string::npos;
|
|
a.stages.clear();
|
|
if (want_tr) a.stages.push_back("train");
|
|
if (want_xp) a.stages.push_back("export");
|
|
if (a.stages.empty()) {
|
|
fprintf(stderr, "ace-train train-dit: --stages must name at least one of train,export\n");
|
|
return 2;
|
|
}
|
|
}
|
|
|
|
// ── numeric sanity (the server clamps these too; be defensive) ───────
|
|
if (a.rank < 1 || a.rank > 256 || a.alpha < 1 || a.alpha > 1024 || a.epochs < 1 || a.epochs > 2000 ||
|
|
a.grad_accum < 1 || a.grad_accum > 64 || a.lr <= 0.0f || a.lr > 1.0f || a.grad_clip < 0.0f ||
|
|
a.grad_clip > 100.0f || a.warmup_ratio < 0.0f || a.warmup_ratio > 0.5f || a.weight_decay < 0.0f ||
|
|
a.weight_decay > 1.0f || a.target_loss < 0.0f || a.target_loss > 20.0f || a.milestone_keep < 0 ||
|
|
a.batch < 1 || a.batch > 16 || a.ckpt < 0 || a.ckpt > 32 ||
|
|
a.milestone_keep > 64 || a.milestone_step < 0.0f || a.milestone_step > 5.0f || a.vram_reserve_mb < 0 ||
|
|
a.vram_reserve_mb > 16384 || a.vram_safety < 0.0f || a.vram_safety >= 1.0f || a.layers < 0 || a.layers > 64 ||
|
|
(a.crop != 0 && (a.crop < 128 || a.crop > 8192)) || a.crop_min < 128 || a.crop_min > 8192 ||
|
|
a.crop_max < a.crop_min || a.crop_max > 8192 || a.snr_gamma < 1.0f || a.snr_gamma > 100.0f ||
|
|
a.t_bias < 0.0f || a.t_bias > 4.0f || a.timestep_mu < -4.0f || a.timestep_mu > 4.0f ||
|
|
a.timestep_sigma <= 0.0f || a.timestep_sigma > 4.0f || a.t_min < 0.0f || a.t_max > 1.0f ||
|
|
a.t_min >= a.t_max || a.cfg_ratio < 0.0f || a.cfg_ratio > 1.0f || a.genre_ratio < 0 || a.genre_ratio > 100) {
|
|
fprintf(stderr, "ace-train train-dit: numeric option out of range\n");
|
|
return 2;
|
|
}
|
|
if (a.order != "shuffle" && a.order != "fixed") {
|
|
fprintf(stderr, "ace-train train-dit: --order must be shuffle|fixed\n");
|
|
return 2;
|
|
}
|
|
// "replace" never applies the tag during preprocessing (preprocess-run.h:203),
|
|
// so an adapter trained from it carried no trigger and must not claim one.
|
|
if (!a.trigger_position.empty() && a.trigger_position != "prepend" && a.trigger_position != "append") {
|
|
fprintf(stderr, "ace-train train-dit: --trigger-position must be prepend|append\n");
|
|
return 2;
|
|
}
|
|
if (a.trigger.size() > 128) {
|
|
fprintf(stderr, "ace-train train-dit: --trigger must be 128 characters or fewer\n");
|
|
return 2;
|
|
}
|
|
if (a.loss_weighting != "none" && a.loss_weighting != "flow_snr") {
|
|
fprintf(stderr, "ace-train train-dit: --loss-weighting must be none|flow_snr\n");
|
|
return 2;
|
|
}
|
|
// Validated HERE (§2.1): dit_train_main keeps a belt-and-braces check, but a
|
|
// typo must not cost a model load first.
|
|
if (a.adapter_type != "lora" && a.adapter_type != "lokr") {
|
|
fprintf(stderr, "ace-train train-dit: --adapter-type must be lora|lokr\n");
|
|
return 2;
|
|
}
|
|
if (a.lokr_dim < 4 || a.lokr_dim > 4096 || a.lokr_alpha < 0.0f || a.lokr_alpha > 8192.0f ||
|
|
(a.lokr_factor != -1 && (a.lokr_factor < 2 || a.lokr_factor > 64))) {
|
|
fprintf(stderr, "ace-train train-dit: --lokr-dim 4-4096, --lokr-alpha 0-8192, --lokr-factor -1 or 2-64\n");
|
|
return 2;
|
|
}
|
|
if (a.mirror != "f32" && a.mirror != "bf16") {
|
|
fprintf(stderr, "ace-train train-dit: --mirror must be f32|bf16\n");
|
|
return 2;
|
|
}
|
|
if (!bwd_valid(a.bwd)) {
|
|
fprintf(stderr, "ace-train train-dit: --bwd must be outprod|mm\n");
|
|
return 2;
|
|
}
|
|
// LoKR's est-vs-peak gap measured ~13 % (the kron-matvec intermediates the
|
|
// fitted arena polynomial never saw — see dit-vram.h's K10 note), so 5 % is
|
|
// not enough margin for it. An explicit --vram-safety always wins.
|
|
if (!safety_user && a.adapter_type == "lokr") {
|
|
a.vram_safety = 0.12f;
|
|
}
|
|
|
|
// ── required arguments ──────────────────────────────────────────────
|
|
if (a.tensors_dir.empty()) {
|
|
fprintf(stderr, "ace-train train-dit: --tensors is required\n");
|
|
return 2;
|
|
}
|
|
if (!a.self_test && a.out_dir.empty()) {
|
|
fprintf(stderr, "ace-train train-dit: --out is required\n");
|
|
return 2;
|
|
}
|
|
|
|
// ── model resolution (same rule as preprocess/train-lm) ─────────────
|
|
// The DiT used for training MUST be the one the latents were preprocessed
|
|
// against — the encoder states and context latents are that model's outputs.
|
|
if (dit_arg.empty()) {
|
|
dit_arg = dit_meta_dit_path(a.tensors_dir.c_str());
|
|
if (dit_arg.empty()) {
|
|
fprintf(stderr, "ace-train train-dit: no --dit and no dit_path in %s/preprocess_meta.json\n",
|
|
a.tensors_dir.c_str());
|
|
return 2;
|
|
}
|
|
}
|
|
{
|
|
ModelRegistry reg;
|
|
if (!a.models_dir.empty()) {
|
|
registry_scan(®, a.models_dir.c_str());
|
|
}
|
|
std::string name;
|
|
if (!resolve_model(dit_arg, reg.dit, false, a.dit_path, name)) {
|
|
fprintf(stderr, "ace-train train-dit: cannot resolve --dit '%s' in %s\n", dit_arg.c_str(),
|
|
a.models_dir.c_str());
|
|
return 2;
|
|
}
|
|
a.dit_name = name;
|
|
}
|
|
|
|
// MUST precede every model load and graph build — the ggml patch latches
|
|
// GGML_BACKWARD_MM on the first backward it emits.
|
|
bwd_apply(a.bwd);
|
|
|
|
// MANDATORY: ggml_time_ms() divides by an uninitialised frequency otherwise.
|
|
ggml_time_init();
|
|
return dit_train_main(a);
|
|
}
|
|
|
|
// ─── main ───────────────────────────────────────────────────────────────────
|
|
|
|
int main(int argc, char ** argv) {
|
|
if (argc < 2) {
|
|
print_usage();
|
|
return 2;
|
|
}
|
|
if (!strcmp(argv[1], "-h") || !strcmp(argv[1], "--help")) {
|
|
print_usage();
|
|
return 0;
|
|
}
|
|
if (!strcmp(argv[1], "preprocess")) {
|
|
return cmd_preprocess(argc - 1, argv + 1);
|
|
}
|
|
if (!strcmp(argv[1], "train-lm")) {
|
|
return cmd_train_lm(argc - 1, argv + 1);
|
|
}
|
|
if (!strcmp(argv[1], "train-dit")) {
|
|
return cmd_train_dit(argc - 1, argv + 1);
|
|
}
|
|
if (!strcmp(argv[1], "spike")) {
|
|
return cmd_spike(argc - 1, argv + 1);
|
|
}
|
|
fprintf(stderr, "ace-train: unknown subcommand '%s'\n", argv[1]);
|
|
print_usage();
|
|
return 2;
|
|
}
|