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

Abierto
  • ¿Algo no está bien?
#31 1 comentario 0 reacciones 0 asignados Ver en GitHub

Los mantenedores suelen responder en 7 días

Nadie ha tomado este issue todavía.

Evaluación

Dificultad
4/5
Tiempo estimado
3-5 días
Aptitud para principiantes
43/100
Tipo de issue
Nueva funcionalidad
Claridad
Bastante claro
Estado de actividad
Activo
Stack tecnológico
python, pytorch
Área
machine-learning

Línea de trabajo

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.

Escrito por el modelo de indexación a partir del texto del issue.

Descripción

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?
Lenguaje dominante
Python
Estrellas
2.6k
Forks
151
Merge medio
9 d 1 h
PR fusionados (30 d)
1

Preparar el entorno

Este proyecto no incluye contenedor de desarrollo, Dockerfile ni guía de contribución, así que la configuración corre por tu cuenta: empieza por su README y consulta nuestra guía para la primera contribución para los pasos generales.

Primeros pasos

  1. Lee el issue completo y luego la guía de contribución del proyecto.
  2. Comenta en el issue que vas a ocuparte — evita que dos personas hagan lo mismo.
  3. Haz un fork del repositorio y trabaja en una rama.
  4. Abre un pull request que haga referencia al número del issue.

Más de OpenMOSS/MOSS-Image

Todos los issues de OpenMOSS/MOSS-Image

Issues similares

Más issues de Python

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.