File size: 5,277 Bytes
5de8603 | 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 | #!/usr/bin/env python3
"""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()
|