Repository navigation
moss_image: the transformer gated block by block against float32, on Metal and CUDA - #2294
Conversation
|
Thanks @BobClawblaw. The block gate reads well: the float32 stream comes from the pinned publisher pipeline, each
|
|
Thanks @ashhart. Items 3 and 4 are done locally: the CHANGELOG line as you worded it, and the body will say "Part of #1529". I'll push them with the rebase onto #2291 (item 1) once it lands. On item 2: I ran #2291's head (0dc4613) on Metal on both doc cases. Each case has float32 values for all 50 steps on both branches, so coverage is complete. I ran it twice: once on main + #2291, and once on this branch rebased onto #2291. The two builds give the same results: byte-identical final latents, and the same worst element at every step. So this PR's forward refactor doesn't change anything #2291 measures. Both builds fail every one of the 100 branch-steps in each case:
At the sticky-note step you asked about (step 23, negative branch), our worst element is 284 times its bound: |diff| 0.0173 against a bound of 6.1e-5. I think this comes from the bound, not from our forward. The same check on torch's own bf16 run of that step on the CPU, scored against the MPS teacher's per-element bound, fails too:
The bound at each element is twice the MPS bf16 run's own error at that element, so it inherits wherever MPS happened to round close to float32. Across a 42-block chain, any other bf16 run lands elsewhere at a sizeable share of the elements, torch's CPU run included. As written, #2291 fails every real MOSS-Image case, on any device. The planted case in its host tests passes because its synthetic bf16 values sit inside the envelope by construction. Ways to gate whole steps against float32 that torch's own bf16 runs would pass:
I'd rather you choose the form. Once #2291 settles, I'll rebase onto it and describe both gates as they stand, with this table in the doc. |
|
Thanks @BobClawblaw, and for the table. #2291 landed this morning, so please push the rebase onto main now, with the CHANGELOG line and "Part of #1529". Describe #2291's gate as it stands, with your table in the doc. Your table makes a fair point: at step 23 the branch is 284× the per-element bound on Metal, and torch's own CPU bf16 run is 353×. We'll settle the bound's form separately, and this PR doesn't need to wait for it. |
b0733fc to
39ea898
Compare
|
Thanks @ashhart. The rebase onto main is pushed as The docs now describe #2291's gate as it landed. One correction to my last comment: the medians are 509 (fox) and 435 (sticky note), not 511 and 444. I had taken the upper middle value of the 100 instead of the median. The min and max are unchanged. On this head the block gates still pass 305/305 on Metal and on the 5090. The required checks and the lean check pass; the body lists where each one ran. |
|
/atc exchange 892c67d43fd79e570efeb4c30f1a1dd99fd66fde18ea5060338a6b4ee332f1ef |
|
/atc exchange e335b80a12fb288f57914b77362b5bf04e7f85f2e919d07d945e9e4f81bd4aff |
39ea898 to
6eb8ad9
Compare
|
Rebased onto main |
|
Thanks @BobClawblaw, the rebase is what we asked for. Both gates are described as they stand, the step-23 table is in the doc, and the body says "Part of #1529". The block gate is a clear improvement: 305 of 305 at half the limit on both backends, and the planted row caught where it should be. It lands once CI has run, and we'll add the thanks to the CHANGELOG line. The teacher bound's form is still to be settled on #1529, separately from this PR. |
|
Thanks @BobClawblaw. Main reached 600 lines in |
6eb8ad9 to
cad4ab2
Compare
…Metal and CUDA tf-moss-image-blocks runs one step's branch through Model.steps (on both backends): the embedding, each kept block alone on a float32 run's stream rounded to bf16, then the final layer. Each must be within twice torch bf16's distance from float32 in cosine error and worst row, and a second run must repeat the first byte for byte. moss_image_case.py blocks records the float32 stream and torch bf16's steps alone, on the case's device and on the CPU; the limit is the farther of the two (reference.wider), since torch's own bf16 runs disagree at a massive-activation row (block 0 of the fox's step 1, negative branch: worst row 0.00687 on MPS, 0.00960 on the CPU). tf-moss-image's whole forwards are reported beside torch bf16's and not gated, as the encoder's whole prompt is not; a repeat that differs still fails it. Seven steps' branches: 305 of 305 gates pass on Metal (at most 0.50 of the limit) and on an RTX 5090 (at most 0.60), each run repeated byte for byte; a row planted 5% off in block 20 fails at blocks 20 and 21. Model.forward's modulation and embedding moved into helpers that Model.steps shares; both cases' latents are byte-identical before and after on both backends.
…o zig/build/cuda.zig stays under 600 lines With main at 600 lines, this branch's tf-moss-image-blocks took zig/build/cuda.zig to 602. Both MOSS-Image tools move, unchanged, to one file that cuda.zig calls once.
cad4ab2 to
4cbffd2
Compare
|
Thank you @BobClawblaw. This is merged into main with your commits under your own name, so you're listed as a TensorFold contributor. It ships in the next release, credited to you in the changelog and the release notes. |
Part of #1529 (its third item: the teacher check gated against float32, as #1512 does for the encoder).
What changes
tf-moss-image-blocks(Metal and CUDA) runs one step's branch throughModel.steps. That is the embedding, each kept block alone on a float32 run's stream rounded to bf16, then the final layer. Each must be within twice torch bf16's distance from float32, in cosine error and in its worst row, and a second run must repeat the first byte for byte. Otherwise it fails.moss_image_case.py blocksrecords that step's float32 stream and torch bf16's steps alone on it. It runs them twice: on the case's device and on the CPU. The limit is the farther of the two (reference.wider; tested inreference.zig).tf-moss-imagekeeps Tests: the MOSS-Image teacher gate compares each step with the teacher, not only with a repeat (#1529) #2291's teacher gate on each whole forward as it landed, and still reports cosine and worst row beside torch bf16's. The docs describe both gates. On both doc cases the teacher gate fails every step, on Metal and on CUDA, and torch's own bf16 run on the CPU fails it too (see Measured). The bound's form is being settled separately.Model.forwardon both backends moves its modulation and embedding into helpers thatModel.stepsshares. Both doc cases give byte-identical latents before and after, on Metal and on an RTX 5090.No shared code changes. No kernels change.
Why two torch runs set the limit
I first gated each whole forward, as one 42-block chain. Metal failed one branch-step of 100 on the sticky note: a worst row 2.55 times torch MPS's, at row 313. Block by block, every block matched torch. torch's own bf16 run on the CPU of that step has a worst row 1.37 times its MPS run's. So the worst row of a 42-block bf16 chain is not a stable gate, and this PR gates block by block.
Per block, one MPS run as the limit failed CUDA once in 305 gates: block 0 of the fox's step 1, negative branch. There the negative prompt's seventh row becomes a massive activation, with one channel near 16,500. I dumped CUDA's block 0 and replayed each operation from CUDA's own bf16 inputs with the publisher's torch code.
situand both residual adds match torch bit for bit. The norms, q/k norm and rotary differ on about 0.001% of elements. The projections differ by one ulp on about 0.2%, which is accumulation order. Attention differs from torch no more than torch CPU differs from torch CUDA. The whole block in torch bf16 on the CPU lands where CUDA does:At every other gate, torch's two runs are within 0.84 to 1.20 times each other in worst row and 0.95 to 1.05 in cosine error.
Measured (this commit)
Seven steps' branches, made on the Mac with torch 2.14.1 and diffusers 0.42.0 (MPS and CPU): the fox's steps 1 and 25 and the sticky note's step 23, both branches each, and the sticky note's step 50, prompt branch.
The CUDA whole forwards use the CUDA doc's cases, recorded on the GB10. A block case runs the whole transformer in float32, which does not fit an RTX 5090, so the CUDA block cases are the Mac's. The block gate was not run on a GB10. On Metal, #2291's head without this PR gives the same latents and the same worst element at every step.
The teacher gate at the sticky note's step 23, negative branch (in the Metal doc):
Checks
Each exits 0 on this head (
6eb8ad969, on mainc70d67bb2;39ea89817rebased,CHANGELOG.mdthe only conflict, the code commit unchanged bygit range-diff), on an M5 Ultra:zig fmt --check build.zig zigzig build test --test-timeout 10mtest-golden(322/335 and 55/55 equal to server goldens, 0 unexpected)test-kernel-contracttest-server-security test-pull-securitycheck-release-pathsfuzz-cortheondev-tools -Doptimize=DebugOn the RTX 5090,
zig build install -Dsm=120,test-cuda-hostand the block gates exited 0 on39ea89817. Main's CUDA changes since then are GLM's L2 prefetch and its contract records, which this PR's CUDA code doesn't use; they were not rerun on this head.Lean: base 155, head 155, 0 new.
Overlap
#2281 (MOSS-Image serving on CUDA) touches two of the same files, and both conflicts are small:
moss_image/cuda.zig: both addpub const reference.zig/build/cuda.zig: moss image: serve pictures on CUDA (encoder, decoder and host for #1573) #2281 moves the MOSS tools intomoss_tools.zig, so this PR's CUDAtf-moss-image-blockstarget would move there too.Whichever lands second rebases.