$ the-wire · showcase
Torchtitan trains MoE routers in fp32 with bf16x9 on Blackwell
By RepoJournal · Filed · About PyTorch
MoE router forward and backward outputs now stay in fp32, using bf16x9 kernels to keep tensor-core throughput on Blackwell.
Torchtitan's new RouterLinear keeps forward output in fp32 and computes backward dgrad and wgrad in fp32, using a "bf16x9" trick to emulate fp32 precision on Blackwell tensor cores; the acceleration from tensor cores outweighs the extra compute ops despite bf16x9 having more compute than plain fp32 [1]. If you train MoE models on Blackwell, expect improved router precision without the usual throughput penalty.
The MoE load-balancing optimizer hook now scans both `model.layers` and `model.mtp_layers`, so DeepSeek MTP decoder blocks no longer keep stale `expert_bias_E` and `tokens_per_expert_E` during distributed training [2]. Separately, Muse Glimmer's final norm used a zero initial effective scale that zeroed hidden states, logits, and first-backward gradients; its weight now initializes to one while retaining `gain_center=0.0`, preserving the checkpoint parameterization [3].
TorchRL's MicroDuck task system now treats a task as data on a tensorclass with one row per environment, adding sidestep and jump presets, fixed per-env task assignment via `MicroDuckTaskSampler(task_ids=)`, and an optional 2x2 evaluation video [4][5]. A new tutorial walks through building task libraries, pinning tasks, and using a reward registry so one policy handles all tasks [6]. The old-dependency CI job is restored by using a pre-2.2 load-state hook API for DreamerV3 and passing string paths to older torch mmap checkpoint loading [7].
TensorDict 0.14.1 ships with fixes for module metadata, lazy and non-tensor indexing, and frozen tensorclass deserialization, with no breaking changes [8]. In PyTorch, a CPU Inductor bug that raised `ZeroDivisionError` for vectorized integer remainder with padded tails (when eager and TorchScript worked) is fixed [9], and Dynamo type replay registries replace `Any` with `object` for more precise typing [10].
Action items
- → Test MoE router training with the new bf16x9 fp32 path on Blackwell pytorch/torchtitan [plan]
- → Inspect your MoE load-balancing hook for MTP layers pytorch/torchtitan [plan]
- → Upgrade TensorDict to 0.14.1 to pick up correctness fixes pytorch/tensordict [plan]
- → Retest integer remainder models on CPU Inductor after the fix pytorch/pytorch [plan]
References
- [1] [router] Fp32 for backward and use bf16x9 kernels ↗ pytorch/torchtitan
- [2] [MoE] Include MTP layers in load-balancing hook ↗ pytorch/torchtitan
- [3] Fix Muse Glimmer final norm initialization ↗ pytorch/torchtitan
- [4] [Refactor] MicroDuck task library: tensorclass tasks, reward registry, sidestep and jump ↗ pytorch/rl
- [5] [Feature] Optional 2x2 task video at MicroDuck evaluations ↗ pytorch/rl
- [6] [Doc] MicroDuck tutorial: tasks as data, rewards as a registry, one policy for all ↗ pytorch/rl
- [7] [CI] Restore old-dependency compatibility ↗ pytorch/rl
- [8] TensorDict v0.14.1 ↗ pytorch/tensordict
- [9] [Inductor] Keep padded tail divisors non-zero in vectorized integer remainder (#192025) ↗ pytorch/pytorch
- [10] [dynamo] Type replay registries as object (#196054) ↗ pytorch/pytorch
