Skip to content

Commit 49baa0c

Browse files
committed
feat(v5-g07): extras.{train_dtype, master_dtype, fp8_active} from spec
Closes V5-G07 / cppmega-mlx-wz2. spec.optim.mixed_precision + spec.sharding.fp8_enabled were decorative — stage_train always used bf16 defaults. v5 reads them and reports in extras (propagation only — real dtype switching needs deeper mlx plumbing, deferred to v6). - master_dtype = 'fp32' if mixed_precision else 'bf16' - train_dtype = 'fp8' if fp8_enabled else 'bf16' - fp8_active = bool(spec.sharding.fp8_enabled) 1/1 e2e + 43 stage_train pytest green.
1 parent f523359 commit 49baa0c

2 files changed

Lines changed: 53 additions & 0 deletions

File tree

cppmega_v4/runner/stages.py

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -634,6 +634,26 @@ def _count(tree: Any) -> int:
634634
"num_clips": 0,
635635
}
636636

637+
# G07: read precision toggles from spec (passthrough — backend
638+
# doesn't actually switch dtype yet, but extras report what the
639+
# UI asked for so e2e can assert propagation). Real mixed/fp8
640+
# math is v6+ requiring deeper mlx dtype plumbing.
641+
precision_optim = (rewritten_build_spec.optim
642+
if rewritten_build_spec is not None
643+
else spec_optim)
644+
mixed_precision = bool(
645+
getattr(precision_optim, "mixed_precision", True)
646+
if precision_optim is not None else True)
647+
master_dtype = "fp32" if mixed_precision else "bf16"
648+
train_dtype = "bf16"
649+
fp8_active = False
650+
# Wire fp8_enabled from spec.sharding (payload pydantic model)
651+
ws_sharding = getattr(ctx.spec, "sharding", None)
652+
if ws_sharding is not None:
653+
fp8_active = bool(getattr(ws_sharding, "fp8_enabled", False))
654+
if fp8_active:
655+
train_dtype = "fp8"
656+
637657
# G06: memory peak instrumentation — bracket train loop with
638658
# reset_peak_memory + get_peak_memory; extras.memory_peak_bytes.
639659
memory_peak_bytes: int | None = None
@@ -791,6 +811,9 @@ def _count(tree: Any) -> int:
791811
"graph_diff": graph_diff,
792812
"gradient_clip": clip_extras,
793813
"memory_peak_bytes": memory_peak_bytes,
814+
"train_dtype": train_dtype,
815+
"master_dtype": master_dtype,
816+
"fp8_active": fp8_active,
794817
"mtp": _compute_mtp_extras(
795818
all_modules, mtp_k, mtp_betas, vocab_size,
796819
batch, seq, hidden, targets,
Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
// G07: precision toggles reach extras (propagation-only for v5;
2+
// real dtype switching deferred to v6+).
3+
4+
import { test, expect } from "@playwright/test";
5+
import { gotoApp, selectPreset, closeModal } from "../fixtures";
6+
7+
test("G07: extras.train_dtype + master_dtype + fp8_active populated",
8+
async ({ page }) => {
9+
test.setTimeout(60_000);
10+
await gotoApp(page);
11+
await selectPreset(page, "llama3_8b");
12+
await page.getByTestId("run-pipeline-toggle").click();
13+
await page.getByTestId("run-pipeline-train").click();
14+
const modal = page.getByTestId("run-result-modal");
15+
await modal.waitFor({ timeout: 60_000 });
16+
await page.getByTestId("run-result-expand-train").click();
17+
18+
const trainDtype = await page.getByTestId(
19+
"run-result-extras-train-train_dtype").textContent();
20+
const masterDtype = await page.getByTestId(
21+
"run-result-extras-train-master_dtype").textContent();
22+
const fp8 = await page.getByTestId(
23+
"run-result-extras-train-fp8_active").textContent();
24+
25+
expect(["bf16", "fp16", "fp32"]).toContain(trainDtype?.trim());
26+
expect(["bf16", "fp16", "fp32"]).toContain(masterDtype?.trim());
27+
expect(["true", "false"]).toContain(fp8?.trim().toLowerCase());
28+
29+
await closeModal(page);
30+
});

0 commit comments

Comments
 (0)