circlestone-labs

/

Anima

Like 2.58k

Follow

CircleStone Labs 1.13k

Hi. Recently I was obsessed to precision issues and might find easy way to improve prediction/gradient accuracy. Running Anima in bf16 rounds the residual stream between blocks, and also the modulation, embedder and head. Keeping those parts in fp32 and leaving only the block linears and attention in bf16 brings the prediction much closer to a full fp32 forward. No weights change dtype, so it costs almost no VRAM and no conversion.

+class Fp32Linear(nn.Linear):
+    def forward(self, x):
+        with torch.autocast(x.device.type, enabled=False):
+            bias = None if self.bias is None else self.bias.float()
+            return F.linear(x.float(), self.weight.float(), bias)

 class RMSNorm(nn.Module):
     def forward(self, x):
         output = self._norm(x.float())
-        return (output * self.weight).to(x.dtype)
+        return output * self.weight.float()  # RoPE then rotates q / k in fp32

+ def stream_fp32_(dit):
+    # x_embedder, final_layer.linear, t_embedder, all adaln_* linears
+    for m in fp32_linears(dit):
+        m.__class__ = Fp32Linear
+    forward = dit.forward_mini_train_dit
+    def under_autocast(x, *args, **kwargs):
+        with torch.autocast(x.device.type, dtype=torch.bfloat16):
+            return forward(x, *args, **kwargs)
+    dit.forward_mini_train_dit = under_autocast

 # sampler
-latents = latents.to(torch.bfloat16)
+latents = latents.to(torch.float32)

Bf16 autocast rounds each block linear's input, and q / k are rounded back to bf16 on entry to FlashAttention. The sampler state stays fp32, which also removes its per-step rounding.

Bench

Prediction error ‖v − v_fp32‖ / ‖v_fp32‖ in %, mean over six images, no LoRA:

One 8-step trajectory at 768², two images. In my personal 5070Ti graphic card:

Though this would require more thorough benchmarks, but doesn't really cost both for integrations and computations.