Hacktoberfest 2026：メンテナが10月に向けて印を付けた、オープンで初心者向けの issue。 Hacktoberfest の issue を見る

perf: Triton normalization backward underperforms apex at wide hidden dimensions

オープン
  • 何か問題がありますか？
#499 コメント 0 件 リアクション 0 件 担当者 0 名 GitHub で見る

まだ誰も着手していません。

評価

難易度
4/5
見積もり時間
3〜5日
初心者へのやさしさ
25/100
issue の種類
リファクタリング
明瞭さ
おおむね明確
活発さ
静か
技術スタック
python
領域
machine-learning, performance

調査の方向性

Start in fast_llm/functional/triton/normalization.py, reading triton_normalization_backward_kernel_1 and triton_normalization_backward_kernel_2, and run tools/benchmark/bench_normalization.py to reproduce the shape table. Phase 1 is a @triton.autotune sweep over kernel_2 tile sizes; Phase 2 restructures kernel_1 into column-chunked reduction and output passes. Done means the benchmark shows gains at the 16K×2K and 2K×16K shapes while its built-in correctness check still passes.

索引モデルが issue の本文から書いたものです。

説明

Summary

The Triton `layer_norm` / `rms_norm` backward is inconsistent across hidden-dimension sizes: it beats `apex_fast` / `apex_fused` at many shapes but falls behind at others. Performance does not degrade monotonically with hidden size — the non-monotonic pattern is explained by the root cause below.

Benchmark: single GPU, H100 SXM, bf16, fwd+bwd. All shapes are isovolumetric (32M elements, varying aspect ratio).

`layer_norm` — Triton vs. apex_fast (fwd+bwd, µs)
shape (rows × cols) Triton apex_fast ratio
32768 × 1024 309 µs (26% BW) 578 µs (14% BW) 2.0× faster
16384 × 2048 619 µs (13% BW) 446 µs (18% BW) 1.4× slower
8192 × 4096 253 µs (32% BW) 578 µs (14% BW) 2.3× faster
4096 × 8192 286 µs (28% BW) 315 µs (25% BW) ~parity
2048 × 16384 371 µs (22% BW) 294 µs (27% BW) 1.3× slower
`rms_norm` — Triton vs. apex_fused (fwd+bwd, µs)
shape (rows × cols) Triton apex_fused pytorch_eager
16384 × 2048 382 µs (21% BW) 440 µs (18% BW) 271 µs (30% BW)
8192 × 4096 226 µs (35% BW) 355 µs (22% BW) ~parity

At (16K, 2K) even plain `pytorch_eager` beats both Triton and apex_fused — the backward kernel is leaving performance on the table.

Root cause

The backward uses two kernels:

`kernel_1` loads an entire row into registers in a single shot:

```python
block_size_row = MAX_BLOCK_SIZE_BYTES / (4 * n_cols)
```

At wide `n_cols` this collapses to 1–4 rows per block, producing many tiny thread blocks and a large partial-sum buffer (`num_blocks_row × n_cols` floats) that `kernel_2` must reduce from DRAM.

`kernel_2` has its tile sizes hardcoded at `block_size_m=64, block_size_n=8`, tuned from a single (8192, 4096) run on an A100. These tile sizes happen to be near-optimal for that one shape, which is why (8Ki, 4Ki) outperforms apex by 2.3× — masking the kernel_1 collapse at that shape. All other shapes suffer from both kernel_1's poor occupancy and kernel_2's mismatched tile sizes, producing the non-monotonic pattern in the table above.

Tensor-parallel note

Under tensor parallelism the hidden dimension is sharded across TP ranks. At TP=4 a 8K-hidden model presents per-rank shapes of 2K cols — right in the range where this issue matters most. This makes the fix load-bearing for real training runs.

Fix plan (in `fast_llm/functional/triton/normalization.py`)

Phase 1 — Autotune `kernel_2` tile sizes (1–2 hours)

Add `@triton.autotune` over a small grid of `(block_size_m, block_size_n)` values for `triton_normalization_backward_kernel_2`. No structural change. Free improvement for all shapes other than (8K, 4K). A `TODO` comment at line 233 already flags the hardcoded tile sizes.

Phase 2 — Two-pass grid-stride column processing in `kernel_1` (1–2 days)

Decouple `block_size_row` from `n_cols` by processing each row in column chunks within the block:

  1. Reduction pass — iterate over column chunks, accumulate two row-wise scalars per row into shared memory:

    • `c1[row] = mean(weight * grad_output * inv_var)` (bias grad correction)
    • `c2[row] = mean(input_normalized * weight * grad_output * inv_var)` (normalization correction)

  2. Output pass — iterate over column chunks again, compute and store `grad_input` using `c1`, `c2` from shared memory; accumulate `grad_weight` partials in shared memory and flush once at the end.

With this structure `block_size_row` becomes a free parameter (32–64 typical) independent of `n_cols`. For (2K, 16K): `num_blocks_row` drops 1024 → 64; partial-sum buffer shrinks 32 MB → 2 MB. Shared memory overhead: `block_size_row × 2 × 4 bytes` scratch for `c1`/`c2` — 512 bytes for 64 rows.

Phase 3 — Eliminate `kernel_2` for small `num_blocks_row` (optional, ~1 day)

After Phase 2, `num_blocks_row` is small enough that `kernel_2` can be bypassed: atomically accumulate `grad_weight` directly from `kernel_1` when `num_blocks_row ≤ threshold`. Phase 2 alone likely closes most of the gap.

Expected outcome

shape current after Phase 1+2
32768 × 1024 fast (2.0× vs apex) ≈ same
16384 × 2048 slow (1.4× slower than apex) fast
8192 × 4096 fast (2.3× vs apex, kernel_2 tuned for this shape) ≈ same
4096 × 8192 ~parity with apex faster
2048 × 16384 slow (1.3× slower than apex) fast

Files

  • `fast_llm/functional/triton/normalization.py` — `triton_normalization_backward_kernel_1`, `triton_normalization_backward_kernel_2`
  • `tools/benchmark/bench_normalization.py` — benchmark harness (correctness check already built in)

The two-pass grid-stride pattern is standard (Flash Attention softmax backward, fused LayerNorm CUDA kernels).

主要言語
Python
スター
337
フォーク
52
30日以内にマージされた PR はありません

環境構築

はじめの一歩

  1. issue を最後まで読み、次にプロジェクトのコントリビューションガイドを読みます。
  2. 着手することを issue にコメントします — 二人が同じ作業をするのを防げます。
  3. リポジトリをフォークし、ブランチを切って変更します。
  4. issue 番号を参照したプルリクエストを送ります。

ServiceNow/Fast-LLM のほかの issue

ServiceNow/Fast-LLM の issue をすべて見る

似ている issue

Python の issue をもっと見る

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。