Skip to content

Cache the x casts to fp8 and build ragged layout for fp8 gemm once #3407

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
MarioSieg wants to merge 1 commit into
base: main
Choose a base branch
from
+79 −20
Open

Cache the x casts to fp8 and build ragged layout for fp8 gemm once #3407

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
71 changes: 58 additions & 13 deletions src/prime_rl/trainer/models/layers/fp8_grouped_gemm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -102,7 +102,11 @@ def forward(
x: torch.Tensor,
weight: torch.Tensor,
offs: torch.Tensor,
layout: tuple | None = None,
x_fp8_cache: tuple[torch.Tensor, torch.Tensor] | None = None,
) -> torch.Tensor:
if layout is None:
layout = build_grouped_layout(offs, total_m=x.size(0))
(
total_m,
padded_total_m,
Expand All @@ -112,19 +116,22 @@ def forward(
starts_tensor,
actual_ms_tensor,
block_starts_tensor,
) = build_grouped_layout(offs, total_m=x.size(0))
) = layout

use_ue8m0 = ue8m0_for_device(x.device)
x_fp8 = grouped_per_token_cast_to_fp8_triton(
x,
padded_total_m,
block_to_group,
starts_tensor,
actual_ms_tensor,
block_starts_tensor,
use_ue8m0,
GROUP_ALIGNMENT,
)
if x_fp8_cache is not None:
x_fp8 = x_fp8_cache
else:
x_fp8 = grouped_per_token_cast_to_fp8_triton(
x,
padded_total_m,
block_to_group,
starts_tensor,
actual_ms_tensor,
block_starts_tensor,
use_ue8m0,
GROUP_ALIGNMENT,
)
weight_fp8 = grouped_per_block_cast_to_fp8_triton(
weight.transpose(1, 2),
use_ue8m0,
Expand Down Expand Up @@ -236,22 +243,60 @@ def backward(ctx, grad_output: torch.Tensor):
block_starts_tensor,
)

return grad_x, grad_weight, None
# grad_x, grad_weight, offs, layout, x_fp8_cache — the last three never need grad.
return grad_x, grad_weight, None, None, None


def compute_grouped_layout(offs: torch.Tensor, total_m: int):
return build_grouped_layout(offs, total_m=total_m)


def cast_grouped_input_to_fp8(x: torch.Tensor, layout: tuple) -> tuple[torch.Tensor, torch.Tensor]:
(
_total_m,
padded_total_m,
_grouped_layout,
block_to_group,
_ks_tensor,
starts_tensor,
actual_ms_tensor,
block_starts_tensor,
) = layout
use_ue8m0 = ue8m0_for_device(x.device)
return grouped_per_token_cast_to_fp8_triton(
x,
padded_total_m,
block_to_group,
starts_tensor,
actual_ms_tensor,
block_starts_tensor,
use_ue8m0,
GROUP_ALIGNMENT,
)


def grouped_fp8_gemm(
x: torch.Tensor,
weight: torch.Tensor,
offs: torch.Tensor,
layout: tuple | None = None,
x_fp8_cache: tuple[torch.Tensor, torch.Tensor] | None = None,
) -> torch.Tensor:
"""FP8 grouped GEMM, drop-in replacement for torch._grouped_mm.

Args:
x: (M, K) concatenated token activations in bfloat16.
weight: (G, K, N) expert weights in bfloat16.
offs: (G,) int32 cumulative token counts per expert.
layout: optional precomputed ``compute_grouped_layout(offs, total_m)`` result,
to skip rebuilding the ragged-group layout when the caller already has one
for this (offs, total_m) pair.
x_fp8_cache: optional precomputed ``cast_grouped_input_to_fp8(x, layout)``
result, to skip re-casting ``x`` when the caller already cast this exact
tensor against this exact layout (e.g. for a sibling grouped_fp8_gemm call
with a different weight but the same x).

Returns:
(M, N) output tensor in bfloat16.
"""
return _GroupedFP8Gemm.apply(x, weight, offs)
return _GroupedFP8Gemm.apply(x, weight, offs, layout, x_fp8_cache)
28 changes: 21 additions & 7 deletions src/prime_rl/trainer/models/layers/moe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -171,11 +171,23 @@ def _run_experts_grouped_mm_impl(
assert x.dim() == 2

if fp8:
from prime_rl.trainer.models.layers.fp8_grouped_gemm import grouped_fp8_gemm
from prime_rl.trainer.models.layers.fp8_grouped_gemm import (
cast_grouped_input_to_fp8,
compute_grouped_layout,
grouped_fp8_gemm,
)

x_bf16 = x.bfloat16()
layout = compute_grouped_layout(offsets, x_bf16.size(0))
x_fp8_cache = cast_grouped_input_to_fp8(x_bf16, layout)

h = F.silu(grouped_fp8_gemm(x.bfloat16(), w1.bfloat16().transpose(-2, -1), offsets))
h = h * grouped_fp8_gemm(x.bfloat16(), w3.bfloat16().transpose(-2, -1), offsets)
out = grouped_fp8_gemm(h, w2.bfloat16().transpose(-2, -1), offsets).type_as(x)
h = F.silu(
grouped_fp8_gemm(x_bf16, w1.bfloat16().transpose(-2, -1), offsets, layout=layout, x_fp8_cache=x_fp8_cache)
)
h = h * grouped_fp8_gemm(
x_bf16, w3.bfloat16().transpose(-2, -1), offsets, layout=layout, x_fp8_cache=x_fp8_cache
)
out = grouped_fp8_gemm(h, w2.bfloat16().transpose(-2, -1), offsets, layout=layout).type_as(x)
else:
h = F.silu(torch._grouped_mm(x.bfloat16(), w1.bfloat16().transpose(-2, -1), offs=offsets))
h = h * torch._grouped_mm(x.bfloat16(), w3.bfloat16().transpose(-2, -1), offs=offsets)
Expand Down Expand Up @@ -1123,10 +1135,12 @@ def _run_nongated_experts_grouped_mm_impl(
assert x.dim() == 2

if fp8:
from prime_rl.trainer.models.layers.fp8_grouped_gemm import grouped_fp8_gemm
from prime_rl.trainer.models.layers.fp8_grouped_gemm import compute_grouped_layout, grouped_fp8_gemm

h = relu2(grouped_fp8_gemm(x.bfloat16(), w1.bfloat16().transpose(-2, -1), offsets))
out = grouped_fp8_gemm(h, w2.bfloat16().transpose(-2, -1), offsets).type_as(x)
x_bf16 = x.bfloat16()
layout = compute_grouped_layout(offsets, x_bf16.size(0))
h = relu2(grouped_fp8_gemm(x_bf16, w1.bfloat16().transpose(-2, -1), offsets, layout=layout))
out = grouped_fp8_gemm(h, w2.bfloat16().transpose(-2, -1), offsets, layout=layout).type_as(x)
else:
h = relu2(torch._grouped_mm(x.bfloat16(), w1.bfloat16().transpose(-2, -1), offs=offsets))
out = torch._grouped_mm(h, w2.bfloat16().transpose(-2, -1), offs=offsets).type_as(x)
Expand Down