Elastic Diffusion for Wan model

[huvunvidia]

No description provided.

[greptile-apps]

Greptile Summary

This PR introduces the first three files of the Elastification / Flextron elastic diffusion method for the Wan model, porting Megatron-LM's FlextronRouter and analytical parameter-count utility into FastGen, along with FlextronConfig and a comprehensive router test suite.

• router.py implements the Gumbel-softmax router with DP-seeded noise, budget interpolation, and a scaler schedule, and deliberately fixes a Megatron NameError in mlp_forward; however, emb_forward silently drops the normalize_router_logits normalization that mlp_forward applies, the scaler tensor is indexed without bounds checking (crashing if curr_iteration ≥ train_iters), and the constructor crashes with TypeError when budget_list=None (the config default).
• budget_utils.py provides the Wan-specific analytical param count for differentiable budget loss; the ScalarOrTensor type alias claims float support but the het-FFN detection branch calls .shape[0] on a plain Python float, raising AttributeError.
• The test suite is well-structured with good axis coverage but does not exercise the normalize_router_logits path on emb_forward or any None-default config scenario.

Confidence Score: 3/5

The router has three independent crash/correctness paths that will be hit in real training runs.

The emb_forward method silently misroutes the embedding axis whenever normalize_router_logits=True, producing inconsistent routing invisible in logs. The scaler index goes out of bounds on the final training iteration or during eval with a large iteration count. The budget_list=None default lets a misconfigured config reach deep into the router constructor before crashing. All three affect the core training loop.

fastgen/methods/elastification/router.py needs the most attention — all three constructor/forward bugs live there. fastgen/methods/elastification/budget_utils.py needs a one-line type-guard fix.

Important Files Changed

Filename	Overview
fastgen/methods/elastification/router.py	Core FlextronRouter implementation with three bugs: emb_forward ignores normalize_router_logits, scaler tensor lacks bounds checking, and budget_list=None crashes the constructor. Also contains a dead repeat_interleave(repeats=1) no-op in skipping_forward.
fastgen/methods/elastification/budget_utils.py	Analytical Wan parameter-count function; ScalarOrTensor type alias claims float support but the het-FFN branch crashes with AttributeError when a Python float is passed for ffn_hidden_size.
fastgen/methods/elastification/config.py	FlextronConfig attrs dataclass with validation; budget_list defaults to None rather than a required field, enabling a crash path in the router constructor that could be caught here instead.
tests/elastification/test_router.py	Thorough router unit tests covering construction, forward paths, DP-seeded Gumbel, tau decay, gradients, skip, hetero-FFN, and the Megatron NameError fix; all fixtures provide valid configs so the budget_list=None and emb_forward/normalize_router_logits gaps are not exercised.
fastgen/methods/elastification/PLAN.md	Design document describing the Elastification port from Megatron-LM; no executable code, well-written with explicit callouts for known Megatron quirks and scope reductions.
tests/elastification/init.py	Empty test package init file; no issues.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant Caller
    participant Router as FlextronRouter.forward()
    participant BudgetMap as budget_map lookup
    participant MLP as mlp_forward()
    participant EMB as emb_forward()
    participant SKIP as skipping_forward()
    participant Gumbel as _dp_gumbel_softmax()

    Caller->>Router: forward(budget, curr_iteration)
    Router->>Router: "hard_sample = random() > hard_sample_th"
    Router->>Router: "tau = tau_init * tau_decay^curr_iter"

    alt budget in budget_map
        Router->>BudgetMap: one_hot lookup
    else "budget == 1.0 fallback"
        Router->>BudgetMap: one_hot for max(budget_map)
    else interpolated budget
        Router->>Router: bucketize + linear blend of two one-hots
    end

    Router->>MLP: mlp_forward(curr_iter, budget_tensor, ...)
    MLP->>MLP: gate_mlp forward pass
    MLP->>MLP: apply scaler + normalize_router_logits branch
    MLP->>Gumbel: _dp_gumbel_softmax(logits, tau, hard)
    note over Gumbel: save/restore CPU+CUDA RNG state
    Gumbel-->>MLP: soft_probs
    MLP-->>Router: (soft_probs, chosen_int)

    opt "add_skipping=True"
        Router->>SKIP: skipping_forward(curr_iter, budget_tensor, ...)
        SKIP->>Gumbel: _dp_gumbel_softmax(logits, tau, hard)
        Gumbel-->>SKIP: soft_probs
        SKIP->>SKIP: build per-layer 0/1 skip mask
        SKIP-->>Router: (soft_probs, skip_mask)
    end

    Router->>EMB: emb_forward(curr_iter, budget_tensor, ...)
    EMB->>EMB: gate_emb forward pass
    note over EMB: normalize_router_logits ignored here
    EMB->>Gumbel: _dp_gumbel_softmax(logits, tau, hard)
    Gumbel-->>EMB: soft_probs
    EMB-->>Router: (soft_probs, chosen_int)

    Router->>Router: "fwd_pass_count += 1"
    Router-->>Caller: (mlp_out, skip_out or None, emb_out)

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant Caller
    participant Router as FlextronRouter.forward()
    participant BudgetMap as budget_map lookup
    participant MLP as mlp_forward()
    participant EMB as emb_forward()
    participant SKIP as skipping_forward()
    participant Gumbel as _dp_gumbel_softmax()

    Caller->>Router: forward(budget, curr_iteration)
    Router->>Router: "hard_sample = random() > hard_sample_th"
    Router->>Router: "tau = tau_init * tau_decay^curr_iter"

    alt budget in budget_map
        Router->>BudgetMap: one_hot lookup
    else "budget == 1.0 fallback"
        Router->>BudgetMap: one_hot for max(budget_map)
    else interpolated budget
        Router->>Router: bucketize + linear blend of two one-hots
    end

    Router->>MLP: mlp_forward(curr_iter, budget_tensor, ...)
    MLP->>MLP: gate_mlp forward pass
    MLP->>MLP: apply scaler + normalize_router_logits branch
    MLP->>Gumbel: _dp_gumbel_softmax(logits, tau, hard)
    note over Gumbel: save/restore CPU+CUDA RNG state
    Gumbel-->>MLP: soft_probs
    MLP-->>Router: (soft_probs, chosen_int)

    opt "add_skipping=True"
        Router->>SKIP: skipping_forward(curr_iter, budget_tensor, ...)
        SKIP->>Gumbel: _dp_gumbel_softmax(logits, tau, hard)
        Gumbel-->>SKIP: soft_probs
        SKIP->>SKIP: build per-layer 0/1 skip mask
        SKIP-->>Router: (soft_probs, skip_mask)
    end

    Router->>EMB: emb_forward(curr_iter, budget_tensor, ...)
    EMB->>EMB: gate_emb forward pass
    note over EMB: normalize_router_logits ignored here
    EMB->>Gumbel: _dp_gumbel_softmax(logits, tau, hard)
    Gumbel-->>EMB: soft_probs
    EMB-->>Router: (soft_probs, chosen_int)

    Router->>Router: "fwd_pass_count += 1"
    Router-->>Caller: (mlp_out, skip_out or None, emb_out)

Loading

_{Reviews (1): Last reviewed commit: "Rewrite budget_utils.py to mirror Megatr..." | Re-trigger Greptile}

[greptile-apps]

emb_forward silently ignores normalize_router_logits

mlp_forward fully implements the normalize_router_logits config flag — normalizing logits by their std before Gumbel-softmax when enabled. emb_forward has no such branch: it only applies the scaler when self.scaler is not None and passes raw logits otherwise, completely ignoring normalize_router_logits. When normalize_router_logits=True, the MLP axis logits are std-normalized while the EMB axis logits are not, producing an asymmetric routing that diverges from the intended configuration and from the Megatron reference.

[greptile-apps]

Scaler tensor indexed without bounds checking

self.scaler is built by torch.linspace(..., steps=train_iters), giving a tensor of length train_iters. All three axis forward methods (mlp_forward, emb_forward, skipping_forward) index it as self.scaler[curr_iteration] without clamping. If curr_iteration >= train_iters — for example, at the final iteration of training (off-by-one), during fine-tuning from a checkpoint with a longer schedule, or in eval mode with a large iteration counter — this raises an IndexError. The fix is to clamp curr_iteration to [0, len(self.scaler) - 1] before indexing.

[greptile-apps]

budget_list=None causes TypeError at construction

FlextronConfig.budget_list defaults to None (line 56 of config.py). FlextronRouter.__init__ calls len(self.config.budget_list) and iterates over it unconditionally on lines 60–65. Constructing FlextronRouter with a default FlextronConfig() (or any config that omits budget_list) crashes immediately with TypeError: object of type 'NoneType' has no len(). A guard — either in FlextronConfig.__attrs_post_init__ that asserts budget_list is not None, or in FlextronRouter.__init__ — would surface this at config-construction time with a clear message rather than a cryptic stack trace mid-constructor.

[greptile-apps]

torch.repeat_interleave(x, repeats=1, dim=0) returns the input tensor unchanged — repeating each element exactly once is a no-op. This dead code (inherited from Megatron, where it likely served a different purpose) adds a needless allocation on every skipping_forward call.

Suggested change

	router_skip_layer_logits = torch.repeat_interleave(
	router_skip_layer_logits, repeats=1, dim=0
	)
	# repeat_interleave(repeats=1) is a no-op; removed.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[greptile-apps]

float type triggers AttributeError in the het-FFN detection branch

The type alias ScalarOrTensor = Union[int, float, torch.Tensor] advertises that ffn_hidden_size can be a Python float, but the branch at line 69 only handles int and torch.Tensor. If a caller passes a plain float (e.g., ffn_hidden_size=5000.0), isinstance(ffn_hidden_size, int) is False, so execution falls to ffn_hidden_size.shape[0], which raises AttributeError: 'float' object has no attribute 'shape'. Either update the type hint to Union[int, torch.Tensor] to reflect the actual contract, or add isinstance(ffn_hidden_size, (int, float)) in the first branch to handle both scalar cases.