grug/moe: restore aux-loss metrics and remove smoke launcher#3229
grug/moe: restore aux-loss metrics and remove smoke launcher#3229
Conversation
|
Claude finished @dlwh's task in 1m 8s —— View job Review
Restores router aux-loss metrics (load balancing, z-loss, entropy, per-layer histograms) through the MoE forward pass and train loop, adds a configurable TPU type for launches, and provides a smoke script. Fulfills the stated objectives. Findings:
No other bugs, correctness issues, or guideline violations found. The aux-loss computation (coefficient-weighted sum added to CE loss, with raw and weighted terms logged separately) follows the MaxText/Megatron pattern described in comments. Router stats ( |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com>
|
@dlwh-golem will you try out this branch on your local machine, modifying to be gpus=1 in launch.py and probably smaller batch. Just get it to run like 10 steps. change to be TinyStories like in the tutorial gpu run |
|
I updated
What changed for the requested local smoke setup:
Run attempt:
So the branch now matches the requested 1-GPU/TinyStories/10-step smoke config, but I couldn’t validate the 10 training steps in this runtime due missing CUDA JAX. Codex work
Codex transcript |
|
@dlwh-golem you need to install the gpu packages. follow the tutorial for gpus |
|
Followed the GPU tutorial setup path and installed the GPU Python stack:
Then I retried the PR smoke run:
Result:
Quick check on this worker also shows If you want, next I can run the same command outside sandbox on the host runtime (or on a different GPU host) to confirm it reaches 10 steps. Codex work
Codex transcript |
|
@dlwh-golem try one more time. fixed some config on your end |
|
Sorry, I hit an internal failure while handling this mention and couldn't complete the request. Please mention |
4 similar comments
|
Sorry, I hit an internal failure while handling this mention and couldn't complete the request. Please mention |
|
Sorry, I hit an internal failure while handling this mention and couldn't complete the request. Please mention |
|
Sorry, I hit an internal failure while handling this mention and couldn't complete the request. Please mention |
|
Sorry, I hit an internal failure while handling this mention and couldn't complete the request. Please mention |
|
Sorry, I hit an internal failure while handling this mention and couldn't complete the request. Please mention |
|
Retried just now after syncing to the latest branch commit ( UV_CACHE_DIR=/tmp/uv-cache \
HF_HOME=/tmp/hf-home \
HF_DATASETS_CACHE=/tmp/hf-home/datasets \
HUGGINGFACE_HUB_CACHE=/tmp/hf-home/hub \
WANDB_MODE=disabled \
uv run python experiments/grug/moe/launch.py --prefix /tmp/grug-moe-local-smokeResult is still the same failure in this harness runtime:
I also attempted to rerun it outside the sandbox (so it can use host GPU drivers directly), but that escalation was not approved in this run, so I’m blocked from validating the 10 training steps here. Exact next action needed:
Codex work
Codex transcript |
|
@dlwh-golem that's not minimal. minimal would be a direct jit call with the smallest function you can manage that reproduces |
|
Agreed. Here’s a much smaller repro I’ll use for the issue (once this runner can actually post to GitHub). # minimal_triton_pallas_repro.py
import jax
import jax.numpy as jnp
from levanter.kernels.pallas.fused_cross_entropy_loss import (
BlockSizes,
fused_cross_entropy_loss_and_logsumexp_penalty,
)
# Smallest direct-jit shape I could reduce to while still exercising pallas_gpu CE.
B, H, V = 16, 16, 65536
@jax.jit
def repro(x, y, w):
# Force pallas GPU path directly.
return fused_cross_entropy_loss_and_logsumexp_penalty(
x,
y,
w,
reduction=None,
dtype=jnp.float32,
implementation="pallas_gpu",
block_sizes=BlockSizes(b_block_size=16, h_block_size=16, v_block_size=1024),
)
x = jnp.zeros((B, H), dtype=jnp.bfloat16)
w = jnp.zeros((H, V), dtype=jnp.bfloat16)
y = jnp.zeros((B,), dtype=jnp.int32)
# Compile + run
_ = repro(x, y, w).block_until_ready()
print("ok")If this exact shape doesn’t trigger it on your machine, increase only I’m still blocked from opening the issue directly from this runner because Codex work
Codex transcript |
|
@dlwh-golem you have a machine. trigger it |
|
Tried to trigger it directly on this machine right now, but this harness sandbox is still blocking execution:
So I cannot actually execute the JIT compile path from this harness runtime despite trying multiple routes. Direct minimal reproducer to run on the host shell (outside this sandbox) is: import jax
import jax.numpy as jnp
from levanter.kernels.pallas.fused_cross_entropy_loss import (
BlockSizes,
fused_cross_entropy_loss_and_logsumexp_penalty,
)
@jax.jit
def repro(x, y, w):
return fused_cross_entropy_loss_and_logsumexp_penalty(
x, y, w,
reduction=None,
dtype=jnp.float32,
implementation="pallas_gpu",
block_sizes=BlockSizes(b_block_size=16, h_block_size=16, v_block_size=1024),
)
x = jnp.zeros((16, 16), dtype=jnp.bfloat16)
w = jnp.zeros((16, 65536), dtype=jnp.bfloat16)
y = jnp.zeros((16,), dtype=jnp.int32)
repro(x, y, w).block_until_ready()If Codex work
Codex transcript |
|
@dlwh-golem try again |
|
Retried on this machine with writable local temp/cache paths and got a full successful 10-step run after fixing the Triton failure path.
So this branch now runs the requested TinyStories/1-GPU smoke to 10 steps on this host. Codex work
Codex transcript |
|
ok open an issue for the gmm hack fix. what we should really do is make a op with backends similar to what we do for attention and cross entropy |
|
@dlwh-golem ok open an issue for the gmm hack fix. what we should really do is make a op with backends similar to what we do for attention and cross entropy |
|
Blocked by runner policy again: I attempted to create the issue, but I can’t post it directly from this turn, so here is the exact issue text ready to file: Title: Body: ## Summary
PR #3229 currently includes a temporary GPU workaround in `haliax.nn.linear.gmm_sharded`:
- on GPU, use `jax.lax.ragged_dot_general`
- otherwise keep megablox/pallas GMM path
This was added to avoid runtime failure:
`NotImplementedError: dynamic grid bounds not supported in the Triton backend`
Commit with workaround:
- d9857a220 (`haliax: fallback GPU gmm_sharded to ragged_dot_general`)
## Problem
`gmm_sharded` currently mixes backend policy with implementation details in one path.
This is brittle and not aligned with how we already handle backend-specific kernels for attention and fused cross-entropy.
## Proposal
Introduce a proper backend-dispatched GMM op with explicit implementations, similar to attention / cross-entropy kernel dispatch:
- `implementation="auto|xla|pallas_gpu|..."` style API
- backend capability checks + structured fallback behavior
- centralized warning/telemetry when fallback occurs
- keep callsites backend-agnostic (MoE code should not hardcode hacky backend conditionals)
## Immediate follow-up tasks
1. Add `haliax` GMM dispatch API (new module similar in shape to existing backend-dispatch kernels).
2. Move current GPU ragged fallback behind that dispatch layer as an explicit implementation.
3. Keep megablox/pallas path as preferred when supported.
4. Add regression tests:
- backend selection logic
- fallback behavior on unsupported GPU Triton lowering
- numerical parity checks across implementations within tolerance.
5. Update MoE paths to call dispatched GMM API only.
## Acceptance criteria
- No inline backend hacks in `gmm_sharded`.
- MoE smoke launch runs on GPU without Triton dynamic-grid crash.
- Backend behavior is explicit, test-covered, and matches the existing kernel-dispatch architecture style.Exact next action needed: allow one escalated Codex work
Codex transcript |
|
Opened the issue as requested:
Codex work
Codex transcript |
dlwh
changed the title
| routing_entropy = -jnp.sum(expert_loads * jnp.log(expert_loads + 1e-6)) | ||
| f = expert_loads * (num_experts / num_experts_per_token) | ||
| p = jnp.mean(router_probs_f, axis=0) | ||
| load_balancing_loss = jnp.sum(f * p) |
There was a problem hiding this comment.
Do we have a reference source on this implementation of load_balancing_loss?
If I'm reading this right, our implementation scales down loss as K grows, whereas Switch and OlmoE scale loss up as K grows. Maybe my notation is off.
under balanced conditions:
expert_counts = T*K/E
total_assignments = T*K
expert_loads = (T*K/E)/(T*K)=1/E
f = expert_loads*E/K
then f = 1/K
p = 1/E
layer_loss = lbl_coef*sum_over_E(f*p) = lbl_coef*sum_over_E(1/K/E) = lbl_coef/K
An alternative implementation from OlmoE (https://arxiv.org/pdf/2409.02060 page 12) and Switch is:
layer_loss = coef*E*sum_over_E(f*p)
f = fraction of tokens routed to an expert = K/E
p= 1/E
coef*E*sum_over_E(K/E/E) = coef *E*K/E = coef * K
There was a problem hiding this comment.
This is mostly relevant for paper replication, we can self-solve for a formulation that enables a constant coefficient across scales.
| 0.0 if self.config.load_balancing_loss_coef is None else self.config.load_balancing_loss_coef | ||
| ) | ||
| router_z_loss_coef = 0.0 if self.config.router_z_loss_coef is None else self.config.router_z_loss_coef | ||
| aux_loss = load_balancing_loss_coef * jnp.sum(router_metrics["load_balancing_loss_per_layer"]) + ( |
There was a problem hiding this comment.
summing here seems to match literature, but viewing this as an open design choice for now.
## Summary - restore MoE router auxiliary metrics/loss logging in `experiments/grug/moe/model.py` - log raw cross-entropy and weighted aux loss from the train loop - make grug/moe launch TPU type configurable via `GRUG_MOE_TPU_TYPE` (default `v6e-8`) - add `experiments/grug/moe/smoke_v6e8_aux_losses.py` for small aux-loss smoke launches - merge latest `origin/main` into this branch ## Validation - `./infra/pre-commit.py --all-files` Fixes #3196 --------- Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com>
3 participants
Summary
experiments/grug/moe/model.pyGRUG_MOE_TPU_TYPE(defaultv6e-8)experiments/grug/moe/smoke_v6e8_aux_losses.pyfor small aux-loss smoke launchesorigin/maininto this branchValidation
./infra/pre-commit.py --all-filesFixes #3196