MiniMax H3: MLP runs fc1 on all tokens at once (6.3 GB at 15 s 1344x768), OOM on 12 GB; token-chunked MLP gives identical output

Open
  • Is something wrong?
#17,703 0 comments 0 reactions 0 assignees View on GitHub

Maintainers usually reply within 1 day

Nobody has claimed this yet.

Assessment

Difficulty
4/5
Estimated time
3-5 days
Newbie friendliness
45/100
Issue type
Bug
Clarity
Clearly specified
Activity status
Active
Tech stack
python, pytorch
Domain
machine-learning, performance

Research direction

Start with MLP.forward and DiTBlock.forward in comfy/ldm/minimax/model.py; the issue describes the full-token allocation and provides a chunked approach. Check how _run_blocks uses the block before changing the in-place residual update. Done when the 15-second T2V workflow at 1344×768 runs on a 12 GB GPU without the fc1 OOM.

Written by the indexing model from the issue text.

Description

MLP.forward (comfy/ldm/minimax/model.py) runs fc1 on every token at once: comfy.ops.linear_input_act(self.fc2, self.fc1(x), "swiglu"). fc1 outputs ffn * 2 = 28672 features, so a 15 s 1344x768 clip (about 109k tokens) needs a 6.3 GB bf16 buffer for that one matmul, and it fails on 12 GB:

  File "comfy_kitchen/backends/cuda/__init__.py", line 2314, in int8_linear
    out = torch.empty((m, n), dtype=out_dtype, device=x.device)
torch.OutOfMemoryError: Allocation on device 0 would exceed allowed memory. (out of memory)

DiTBlock.forward then also keeps the full-size MLP output and two normalized copies of the stream alive at the same time (h from norm1 stays referenced while norm2's h is built).

Environment: ComfyUI v0.45.0 (code also checked unchanged on master b2f6707443), comfy-aimdo 0.6.0, comfy-kitchen 0.2.41, PyTorch 2.15.0+cu132, RTX 3060 12 GB, Linux, Dynamic VRAM on (default). Workflow: the official video_fastvideo_fasth3_t2v / i2v / video_fasth3_ominiref_r2v templates with length 362 (15 s) at 1344x768, about 109k video tokens.

Fix direction. The MLP acts on each token independently. Running it over token chunks (16384 tokens) and adding each chunk's gated output straight into the residual stream gives the same result with about 0.9 GB peak. Locally, this replacement for DiTBlock.forward matched the original bit for bit (max diff 0.0) on test inputs:

def forward(self, x, t_emb, mod_segments, rope_freqs, transformer_options={}, attention=None):
    attention = self.attn if attention is None else attention
    shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaln_proj(t_emb)
    h = _mod_scale_shift(self.norm1(x), shift_msa, scale_msa, mod_segments)
    x = _mod_gate(x, gate_msa, attention(h, rope_freqs=rope_freqs, transformer_options=transformer_options), mod_segments)
    del h
    h = _mod_scale_shift(self.norm2(x), shift_mlp, scale_mlp, mod_segments)
    for c0 in range(0, x.shape[0], CHUNK):
        c1 = min(x.shape[0], c0 + CHUNK)
        out = self.mlp(h[c0:c1])
        for a, b, row in mod_segments:          # gated add, restricted to this chunk
            lo, hi = max(a, c0), min(b, c1)
            if lo < hi:
                r = row[lo - a:hi - a] if torch.is_tensor(row) and row.dim() > 0 else row
                x[lo:hi].addcmul_(out[lo - c0:hi - c0], _mod_row(gate_mlp, r, x.dtype))
    return x

The in-place addcmul_ assumes x is not needed elsewhere after the block, which holds in _run_blocks. Done when 15 s T2V at 1344x768 runs on 12 GB without the fc1 OOM.

Filed by Claude Opus 5.5 in Claude Code.

Dominant language
Python
Stars
146k
Forks
17.4k
Avg merge
19h 44m
Merged PRs (30d)
178

Getting set up

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from Comfy-Org/ComfyUI

All issues in Comfy-Org/ComfyUI

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.