| |
| """Freeze immutable smoke and production plans after virtual manifests pass.""" |
| from __future__ import annotations |
|
|
| from datetime import datetime |
| import hashlib |
| import json |
| import math |
| from pathlib import Path |
|
|
|
|
| HERE = Path(__file__).resolve().parent |
| ROOT = Path("/e/scratch/reformo/schuhmann1_moss/out/m2_600m_ladder/caption_curriculum_v1") |
| V3 = Path("/e/home/jusers/schuhmann1/jupiter/m2_600m_tts_training/cascade_v3_timed") |
|
|
|
|
| def sha(path): |
| h = hashlib.sha256() |
| with Path(path).open("rb") as stream: |
| for block in iter(lambda: stream.read(8 << 20), b""): |
| h.update(block) |
| return h.hexdigest() |
|
|
|
|
| def stages(global_batch, max_updates=None): |
| result = [] |
| for number in range(1, 11): |
| manifest = ROOT / "manifests" / f"S{number}" / "manifest.json" |
| spec = json.loads(manifest.read_text()) |
| assert spec["status"] == "complete" and spec["format"] == "m2-caption-curriculum-v1" |
| natural = math.ceil(spec["presentations"] / global_batch) |
| updates = natural if max_updates is None else min(max_updates, natural) |
| result.append({"name": f"S{number}", "tier": spec["tier"], "manifest": str(manifest), |
| "manifest_sha256": sha(manifest), "presentations": spec["presentations"], |
| "measured_hours": spec["hours"], "updates": updates, |
| **({"max_updates": updates} if max_updates is not None else {})}) |
| return result |
|
|
|
|
| def base(nodes, samples_per_gpu): |
| world = nodes * 4 |
| code = { |
| "continuous_train.py": sha(HERE / "continuous_train.py"), |
| "large_talker.py": sha(HERE / "large_talker.py"), |
| "caption_curriculum_dataset.py": sha(HERE / "caption_curriculum_dataset.py"), |
| "cascade_dataset.py": sha(V3 / "cascade_dataset.py"), |
| "timed_prompt.py": sha(V3 / "timed_prompt.py"), |
| "anno_timed.py": sha(V3 / "anno_timed.py"), |
| } |
| return { |
| "kind": "m2_continuous_large_talker_v1", |
| "architecture": "qwen3_0.6b_plus_fresh_sft3_width_talker", |
| "initialization": {"semantic": "original_pretrained_Qwen3-0.6B", "talker": "random_fresh"}, |
| "nodes": nodes, "ranks_per_node": 4, "world_size": world, |
| "samples_per_gpu": samples_per_gpu, "global_batch": world * samples_per_gpu, |
| "seed": 20260923, "max_padded_tokens": 12000, "max_examples": 16, |
| "lr_backbone_peak": 8e-5, "lr_talker_peak": 2.4e-4, |
| "lr_floor_factor": 0.05, |
| "scheduler": "single_linear_warmup_then_cosine_no_stage_restarts", |
| "gradient_communication": "bfloat16", "deterministic": False, |
| "checkpoint_policy": "all_stage_boundaries_plus_interval", |
| "checkpoint_interval_seconds": 3600, "max_invocation_seconds": 40000, |
| "resume_policy": "latest_complete_same_plan_world", |
| "prompt_contract": "m2-multiformat-caption-v1-20260923", |
| "allowed_prompt_format_ids": ["m2-cascade-timed-repair-v1-20260920", |
| "m2-multiformat-caption-v1-20260923"], |
| "codec_contract": "MOSS-Audio-Tokenizer-v2; 12 codebooks/frame; 12.5 frames/s; 150 codebook tokens/s", |
| "code_sha256": code, "created_at": datetime.now().astimezone().isoformat(), |
| } |
|
|
|
|
| def write_once(path, payload): |
| assert not path.exists(), f"Refusing to overwrite immutable plan {path}" |
| path.write_text(json.dumps(payload, indent=2) + "\n") |
|
|
|
|
| def main(): |
| one = base(1, 2); one["stages"] = stages(one["global_batch"], 1) |
| one["total_updates"] = sum(x["updates"] for x in one["stages"]); one["warmup_steps"] = 2 |
| one["checkpoint_interval_seconds"] = 600; one["max_invocation_seconds"] = 2400 |
| one["output"] = str(ROOT / "smoke_4gpu"); one["smoke_mode"] = True |
| one["smoke_checkpoint_steps"] = [1, one["total_updates"]] |
| one["purpose"] = "real-container schema, prompt, backward and resume smoke" |
| write_once(HERE / "plan_caption_smoke_4gpu.json", one) |
|
|
| wide = base(32, 32); wide["stages"] = stages(wide["global_batch"], 3) |
| wide["total_updates"] = sum(x["updates"] for x in wide["stages"]); wide["warmup_steps"] = 3 |
| wide["checkpoint_interval_seconds"] = 600; wide["max_invocation_seconds"] = 3600 |
| wide["output"] = str(ROOT / "smoke_32nodes"); wide["smoke_mode"] = True |
| wide["smoke_checkpoint_steps"] = [1, wide["total_updates"]] |
| wide["max_acceptable_median_step_seconds"] = 5.0 |
| wide["purpose"] = "32-node throughput acceptance gate for expanded caption mix" |
| write_once(HERE / "plan_caption_smoke_32nodes.json", wide) |
|
|
| production = base(32, 32); production["stages"] = stages(production["global_batch"]) |
| production["total_updates"] = sum(x["updates"] for x in production["stages"]) |
| production["warmup_steps"] = round(0.05 * production["total_updates"]) |
| production["output"] = str(ROOT / "training") |
| production["purpose"] = "authorized continuous S1-to-S10 multi-caption production" |
| write_once(HERE / "plan_caption_production_32nodes.json", production) |
| print(json.dumps({"smoke_4gpu": one["total_updates"], "smoke_32nodes": wide["total_updates"], |
| "production": production["total_updates"], |
| "warmup": production["warmup_steps"]}, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
