Hacktoberfest 2026: los issues que los mantenedores marcaron para octubre, abiertos y aptos para principiantes. Explorar issues de Hacktoberfest

perf: Triton normalization backward underperforms apex at wide hidden dimensions

Abierto
  • ¿Algo no está bien?
#499 0 comentarios 0 reacciones 0 asignados Ver en GitHub

Nadie ha tomado este issue todavía.

Evaluación

Dificultad
4/5
Tiempo estimado
3-5 días
Aptitud para principiantes
25/100
Tipo de issue
Refactorización
Claridad
Bastante claro
Estado de actividad
Tranquilo
Stack tecnológico
python
Área
machine-learning, performance

Línea de trabajo

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.

Escrito por el modelo de indexación a partir del texto del issue.

Descripción

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

Lenguaje dominante
Python
Estrellas
337
Forks
52
Sin PR fusionados en 30 d

Preparar el entorno

Primeros pasos

  1. Lee el issue completo y luego la guía de contribución del proyecto.
  2. Comenta en el issue que vas a ocuparte — evita que dos personas hagan lo mismo.
  3. Haz un fork del repositorio y trabaja en una rama.
  4. Abre un pull request que haga referencia al número del issue.

Más de ServiceNow/Fast-LLM

Todos los issues de ServiceNow/Fast-LLM

Issues similares

Más issues de Python

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.