Hacktoberfest 2026: những issue maintainer đã đánh dấu cho tháng Mười, đang mở và phù hợp người mới. Xem issue Hacktoberfest

perf: Triton normalization backward underperforms apex at wide hidden dimensions

Đang mở
  • Có gì không ổn?
#499 0 bình luận 0 reaction 0 người được giao Xem trên GitHub

Chưa có ai nhận issue này.

Đánh giá

Độ khó
4/5
Thời gian dự kiến
3-5 ngày
Mức phù hợp với người mới
25/100
Loại issue
Tái cấu trúc
Độ rõ ràng
Khá rõ ràng
Mức độ hoạt động
Ít trao đổi
Công nghệ
python
Lĩnh vực
machine-learning, performance

Hướng nghiên cứu

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.

Do mô hình lập chỉ mục viết ra từ nội dung của issue.

Mô tả

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).

Ngôn ngữ chính
Python
Star
337
Fork
52
Không có pull request nào được merge trong 30 ngày

Chuẩn bị môi trường

Bắt đầu từ đâu

  1. Đọc hết issue, rồi đọc hướng dẫn đóng góp của dự án.
  2. Bình luận trên issue rằng bạn sẽ nhận — tránh hai người làm cùng một việc.
  3. Fork repository và làm thay đổi trên một nhánh.
  4. Mở pull request có tham chiếu số hiệu của issue.

Issue khác của ServiceNow/Fast-LLM

Tất cả issue của ServiceNow/Fast-LLM

Issue tương tự

Thêm issue về Python

Nhận issue mới trong hộp thư của bạn

Bản tóm tắt ngắn những issue GitHub phù hợp với người mới.