Hacktoberfest 2026:维护者为十月标记出来的 issue,仍然开放、适合新手。 浏览 Hacktoberfest issue

Most of the bf16 inference error comes from the bf16 residual stream; an fp32 stream around the bf16 GEMMs costs ~1.5%

未关闭
#31 0 条评论 0 个 reaction 已指派 0 人 在 GitHub 查看

维护者通常 1 天内回复

还没有人认领这个 Issue。

评估

难度
4/5
预计耗时
3-5 天
新手友好度
43/100
Issue 类型
功能
描述清晰度
基本清楚
活跃度
活跃
技术栈
python, pytorch

调研方向

Start by reading MOSSImagePipeline.from_pretrained and the transformer’s _Block and _attention implementations, which the issue describes as handling dtype casts. Check how training sets activation precision and whether the pipeline can expose a stream dtype option. Done means the maintainers resolve the training precision question and decide whether to add an option for fp32 residual streams, with behavior and tests specified.

由索引模型根据 Issue 内容生成。

描述

Hi, and thanks for releasing MOSS-Image. While building LoRA training and inference
tooling for it, we measured where the bf16 numerical error in the DiT comes from. The
result was larger than we expected, so we are sharing it here.

Setting

MOSSImagePipeline.from_pretrained loads the transformer with torch_dtype=torch.bfloat16
and nothing runs under autocast. So the residual stream x, the modulation, the gates and
the norm outputs are all bf16. Each block's x = x + gate * proj(...) rounds to an 8-bit
mantissa twice per block, over 42 blocks.

As a test we kept the 252 block linears (layers.N.proj.{q,k,v,o,ffn_up_gate,ffn_down})
in bf16 and cast every other parameter to fp32. The stream, modulation, gates, norms,
embedders and head then run in fp32 around the same bf16 block GEMMs. The model code
needs no change: _Block already casts its GEMM inputs to self.proj["q"].weight.dtype,
and _attention casts q / k to v.dtype, so attention stays bf16 too. On top of the
README example:

 pipe.scheduler = type(pipe.scheduler).from_config(pipe.scheduler.config)
+for name, p in pipe.transformer.named_parameters():
+    if not (name.startswith("layers.") and ".proj." in name):  # all but the 252 block linears
+        p.data = p.data.float()
 pipe.to("cuda")

Prediction error

MOSS-Image-Base. Six images at ~4.1k image + 184–383 text tokens, σ ∈ {0.1, 0.5, 0.9},
one fixed ε per image. Each run is scored as ‖v − v_fp32‖ / ‖v_fp32‖ against an fp32
forward: every GEMM in fp32 on the same bf16 weights, fp32 norms, SDPA, TF32 off.

σ bf16 (as shipped) fp32 stream bf16, no reduced-precision reduction loss bf16 / fp32 stream / fp32
0.1 4.81% 1.64% 4.80% 0.5635 / 0.5612 / 0.5609
0.5 1.93% 0.81% 1.94% 0.2253 / 0.2249 / 0.2249
0.9 2.08% 0.59% 2.12% 0.3819 / 0.3814 / 0.3814
  • The fp32 stream cuts the prediction error 2.4–3.5x.
  • Its loss matches fp32 to the fourth decimal.
  • Turning off cuBLAS reduced-precision reduction (allow_bf16_reduced_precision_reduction = False)
    does not move it. The error is in the stream, not the GEMM accumulation.

Effect on samples

Same prompts, seeds and initial noise. 1024², 50 steps, CFG 3 with block 22 skipped on
the uncond branch, shift 5. Mean |Δ| per image between bf16 and the fp32 stream, out of
255, over 3 prompts × 2 seeds:

mean Δ range
Base 9.5 5.0–16.4
SFT 7.3 3.5–14.0

Base usually keeps the composition and changes details: eyes, accessories, hands,
small marks. SFT can change the pose. In the third image below, the hands, legs and bag
move. To our eyes neither precision is consistently better; both give clean images.

Three comparisons (bf16 | fp32 stream | |diff| ×4) follow in the first comment: Base
prompt 0 seed 1 (16.4/255), SFT prompt 2 seed 1 (14.0/255) and Base prompt 1 seed 1
(11.8/255).

Image Image Image

These renders use our own sampler setup: an ER-SDE solver, Spectrum feature caching,
an INT8 text encoder and compiled blocks. The prediction-error table above is a single
forward and does not depend on it.

Cost. 20.8 s vs 20.5 s per 1024² image (+1.5%) and about 0.6 GB for the fp32
non-block parameters (RTX 5070 Ti, compiled blocks).

Questions

  1. What precision did the residual stream have in pre-training and SFT? The comment
    "Match the reference activation dtype under fp32 weights + autocast" suggests fp32
    master weights under bf16 autocast. If so, the linear outputs and the stream were
    probably bf16 there too. The shipped bf16 inference would then match training, and the
    fp32 stream would only be closer to exact arithmetic. If the stream was fp32 in
    training, the fp32 stream would also match training better.
  2. If that is useful, would you consider an option in the pipeline (e.g.
    stream_dtype=torch.float32) for users who want the more exact forward at ~1.5% cost?
主要语言
Python
星标
2.2k
派生
114
平均合并
9 天 1 小时
30 天内合并 PR
1

环境准备

这个项目没有提供开发容器、Dockerfile 或贡献指南,环境需要你自己搭建:先看它的 README,通用步骤见我们的新手贡献指南。

从这里开始

  1. 先读完整个 Issue,再读项目的贡献指南。
  2. 在 Issue 下留言说明你要接手 —— 这能避免两个人做同样的事。
  3. Fork 仓库,在一个分支上完成修改。
  4. 提交 Pull Request,并在描述里引用这个 Issue 编号。

OpenMOSS/MOSS-Image 的其他 Issue

查看 OpenMOSS/MOSS-Image 的全部 Issue

相似的 Issue

更多 Python Issue

把新 issue 发到你的邮箱

精选适合新手参与的 GitHub issue 摘要。