Skip to content

feat(router/gemm): add sigmoid and fused dw accumulation for megatron #1

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
hscos wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat(router/gemm): add sigmoid and fused dw accumulation for megatron #1

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
182 changes: 142 additions & 40 deletions sonicmoe/functional/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,24 +6,78 @@

import torch
import torch.nn.functional as F
from quack.gemm_interface import gemm, gemm_dgated, gemm_gated
from quack.gemm_interface import gemm, gemm_add_inplace, gemm_dgated, gemm_gated

from ..enums import ActivationType, is_glu
from .backward import (
_down_projection_backward_act,
_token_broadcast_backward,
_sigmoid_over_topk_bwd,
_topk_over_sigmoid_bwd,
_topk_softmax_bwd,
_up_projection_backward_act,
)
from .forward import _router_forward, _topk_softmax_fwd
from .triton_kernels import TC_topk_router_metadata_triton, general_routing_router_metadata_triton


def _base_parameter_from_view(tensor: torch.Tensor) -> torch.nn.Parameter | None:
base = getattr(tensor, "_base", None)
if isinstance(base, torch.nn.Parameter):
return base
if isinstance(tensor, torch.nn.Parameter):
return tensor
return None


def _megatron_main_grad_buffer(weight_view: torch.Tensor) -> torch.Tensor | None:
if torch.compiler.is_compiling():
return None
base = _base_parameter_from_view(weight_view)
if base is None:
return None
if not hasattr(base, "grad_added_to_main_grad"):
return None
main_grad = getattr(base, "main_grad", None)
if not isinstance(main_grad, torch.Tensor):
return None
if main_grad.shape != base.shape or main_grad.device != base.device:
return None
return main_grad


def _megatron_wgrad_return(weight_view: torch.Tensor) -> torch.Tensor | None:
base = _base_parameter_from_view(weight_view)
if base is None:
return None
if not hasattr(base, "grad_added_to_main_grad"):
return None
base.grad_added_to_main_grad = True

# Megatron's DDP hook needs param.grad to be non-None to mark the bucket
# ready, but the real dW has already been accumulated into main_grad.
dummy = (
torch.zeros((), dtype=base.dtype, device=base.device)
if getattr(base, "zero_out_wgrad", False)
else torch.empty((), dtype=base.dtype, device=base.device)
)
base.grad = dummy.expand_as(base)
return None


class TC_Softmax_Topk_Router_Function(torch.autograd.Function):
@staticmethod
def forward(
ctx, router_logits: torch.Tensor, E: int, K: int, is_softmax_over_topk: bool, norm_topk_probs: bool
ctx,
router_logits: torch.Tensor,
E: int,
K: int,
is_softmax_over_topk: bool,
norm_topk_probs: bool,
router_score_function: str,
) -> tuple[torch.Tensor, torch.Tensor]:
if router_score_function not in {"softmax", "sigmoid"}:
raise ValueError(f"unexpected router_score_function ({router_score_function})")
T = router_logits.size(0)

topk_router_score = torch.empty(T, K, dtype=torch.float32, device=router_logits.device)
Expand All @@ -37,6 +91,7 @@ def forward(
K,
is_softmax_over_topk=is_softmax_over_topk,
norm_topk_probs=norm_topk_probs,
is_sigmoid_router=(router_score_function == "sigmoid"),
)

# Save router_logits for topk(softmax()) backward (recompute full softmax).
Expand All @@ -46,6 +101,7 @@ def forward(
ctx.dtype = router_logits.dtype
ctx.is_softmax_over_topk = is_softmax_over_topk
ctx.norm_topk_probs = norm_topk_probs
ctx.router_score_function = router_score_function

return topk_router_score, topk_router_indices

Expand All @@ -56,20 +112,24 @@ def backward(ctx, dtopk_score: torch.Tensor, _: torch.Tensor):
topk_router_score, topk_router_indices, router_logits = ctx.saved_tensors
dlogits = torch.zeros(T, ctx.E, dtype=ctx.dtype, device=topk_router_score.device)

_topk_softmax_bwd(
router_logits,
dlogits,
None,
dtopk_score,
topk_router_score,
topk_router_indices,
E,
K,
is_softmax_over_topk=ctx.is_softmax_over_topk,
norm_topk_probs=ctx.norm_topk_probs,
)

return dlogits, None, None, None, None
if ctx.router_score_function == "sigmoid":
sigmoid_bwd = _sigmoid_over_topk_bwd if ctx.is_softmax_over_topk else _topk_over_sigmoid_bwd
sigmoid_bwd(dlogits, dtopk_score, topk_router_score, topk_router_indices, K)
else:
_topk_softmax_bwd(
router_logits,
dlogits,
None,
dtopk_score,
topk_router_score,
topk_router_indices,
E,
K,
is_softmax_over_topk=ctx.is_softmax_over_topk,
norm_topk_probs=ctx.norm_topk_probs,
)

return dlogits, None, None, None, None, None


class _UpProjection(torch.autograd.Function):
Expand All @@ -90,6 +150,7 @@ def forward(
activation_type: ActivationType,
is_inference_mode_enabled: bool,
concat_layout: bool = False,
accumulate_wgrad_into_main_grad: bool = False,
) -> torch.Tensor:
T, H = x.shape
I, H, E = w1.shape
Expand Down Expand Up @@ -131,6 +192,7 @@ def forward(
ctx.is_each_token_has_variable_activated_experts = is_each_token_has_variable_activated_experts
ctx.is_glu_activation = is_glu_activation
ctx.concat_layout = concat_layout
ctx.accumulate_wgrad_into_main_grad = accumulate_wgrad_into_main_grad

ctx.save_for_backward(
x,
Expand Down Expand Up @@ -158,6 +220,7 @@ def backward(ctx, _: None, dh: torch.Tensor):
is_glu_activation = ctx.is_glu_activation
is_each_token_has_variable_activated_experts = ctx.is_each_token_has_variable_activated_experts
concat_layout = ctx.concat_layout
accumulate_wgrad_into_main_grad = ctx.accumulate_wgrad_into_main_grad

(
x,
Expand All @@ -171,7 +234,8 @@ def backward(ctx, _: None, dh: torch.Tensor):
) = ctx.saved_tensors

dx_expanded = torch.empty(TK, H, dtype=dh.dtype, device=dh.device)
dw1 = torch.empty_like(w1)
dw1_main_grad = _megatron_main_grad_buffer(w1) if accumulate_wgrad_into_main_grad else None
dw1 = None if dw1_main_grad is not None else torch.empty_like(w1)
db1 = None if b1 is None else torch.empty_like(b1)

_up_projection_backward_act(
Expand All @@ -184,16 +248,30 @@ def backward(ctx, _: None, dh: torch.Tensor):
concat_layout=concat_layout,
)

gemm(
x.T,
dh,
out=dw1.permute(2, 1, 0),
cu_seqlens_k=expert_frequency_offset,
A_idx=x_gather_idx,
batch_idx_permute=None,
dynamic_scheduler=False,
concat_layout=(("out",) if concat_layout else None),
)
if dw1_main_grad is None:
gemm(
x.T,
dh,
out=dw1.permute(2, 1, 0),
cu_seqlens_k=expert_frequency_offset,
A_idx=x_gather_idx,
batch_idx_permute=None,
dynamic_scheduler=False,
concat_layout=(("out",) if concat_layout else None),
)
dw1_return = dw1
else:
gemm_add_inplace(
x.T,
dh,
dw1_main_grad.permute(0, 2, 1),
cu_seqlens_k=expert_frequency_offset,
A_idx=x_gather_idx,
batch_idx_permute=None,
dynamic_scheduler=False,
concat_layout=(("out",) if concat_layout else None),
)
dw1_return = _megatron_wgrad_return(w1)

dx_reduced = torch.empty(T, H, dtype=dh.dtype, device=dh.device)

Expand All @@ -207,7 +285,7 @@ def backward(ctx, _: None, dh: torch.Tensor):
is_varlen_K=is_each_token_has_variable_activated_experts,
)

return dx_reduced, dw1, db1, *[None] * 13
return dx_reduced, dw1_return, db1, *[None] * 13


class _DownProjection(torch.autograd.Function):
Expand All @@ -228,6 +306,7 @@ def forward(
num_activated_expert_per_token_offset: torch.Tensor,
is_varlen_K: bool,
activation_type: ActivationType,
accumulate_wgrad_into_main_grad: bool = False,
) -> torch.Tensor:
TK = a.size(0)
H, I, E = w2.shape
Expand All @@ -254,6 +333,7 @@ def forward(
ctx.K = K
ctx.is_varlen_K = is_varlen_K
ctx.activation_type = activation_type
ctx.accumulate_wgrad_into_main_grad = accumulate_wgrad_into_main_grad

ctx.save_for_backward(
h,
Expand All @@ -273,6 +353,7 @@ def backward(ctx, dout: torch.Tensor):
K = ctx.K
is_varlen_K = ctx.is_varlen_K
activation_type = ctx.activation_type
accumulate_wgrad_into_main_grad = ctx.accumulate_wgrad_into_main_grad

(
h,
Expand All @@ -284,7 +365,8 @@ def backward(ctx, dout: torch.Tensor):
s_scatter_idx,
) = ctx.saved_tensors

dw2 = torch.empty_like(w2)
dw2_main_grad = _megatron_main_grad_buffer(w2) if accumulate_wgrad_into_main_grad else None
dw2 = None if dw2_main_grad is not None else torch.empty_like(w2)
db2 = None if b2 is None else torch.empty_like(b2)
dh = torch.empty_like(h)

Expand All @@ -310,21 +392,34 @@ def backward(ctx, dout: torch.Tensor):
activation_type=activation_type.value,
)

gemm(
dout.T,
a_prime,
out=dw2.permute(2, 0, 1),
cu_seqlens_k=expert_frequency_offset,
A_idx=x_gather_idx,
batch_idx_permute=None,
dynamic_scheduler=False,
)
if dw2_main_grad is None:
gemm(
dout.T,
a_prime,
out=dw2.permute(2, 0, 1),
cu_seqlens_k=expert_frequency_offset,
A_idx=x_gather_idx,
batch_idx_permute=None,
dynamic_scheduler=False,
)
dw2_return = dw2
else:
gemm_add_inplace(
dout.T,
a_prime,
dw2_main_grad,
cu_seqlens_k=expert_frequency_offset,
A_idx=x_gather_idx,
batch_idx_permute=None,
dynamic_scheduler=False,
)
dw2_return = _megatron_wgrad_return(w2)

# TC top-K routing
if not is_varlen_K:
ds = ds.view(T, K)

return None, dh, dw2, db2, ds, *[None] * 10
return None, dh, dw2_return, db2, ds, *[None] * 10


def moe_TC_softmax_topk_layer(
Expand All @@ -340,15 +435,17 @@ def moe_TC_softmax_topk_layer(
is_inference_mode_enabled: bool = False,
is_softmax_over_topk: bool = True,
norm_topk_probs: bool = False,
router_score_function: str = "softmax",
concat_layout: bool = False,
accumulate_wgrad_into_main_grad: bool = False,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
assert ((b1 is None) and (b2 is None)) or (
(b1 is not None) and (b2 is not None)
), "b1 and b2 has to be None or not None at the same time!"
E = router_w.size(0)
router_logits = F.linear(x, router_w)
topk_scores, topk_indices = TC_Softmax_Topk_Router_Function.apply(
router_logits, E, K, is_softmax_over_topk, norm_topk_probs
router_logits, E, K, is_softmax_over_topk, norm_topk_probs, router_score_function
)

T, K = topk_indices.size()
Expand Down Expand Up @@ -386,6 +483,7 @@ def moe_TC_softmax_topk_layer(
activation_type,
is_inference_mode_enabled,
concat_layout,
accumulate_wgrad_into_main_grad,
)

o = _DownProjection.apply(
Expand All @@ -403,6 +501,7 @@ def moe_TC_softmax_topk_layer(
None,
False, # is_each_token_has_variable_activated_expert
activation_type,
accumulate_wgrad_into_main_grad,
)

return o, router_logits, expert_frequency
Expand Down Expand Up @@ -433,6 +532,7 @@ def moe_general_routing_inputs(
activation_type: ActivationType,
is_inference_mode_enabled: bool = False,
concat_layout: bool = False,
accumulate_wgrad_into_main_grad: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
assert ((b1 is None) and (b2 is None)) or (
(b1 is not None) and (b2 is not None)
Expand Down Expand Up @@ -484,6 +584,7 @@ def moe_general_routing_inputs(
activation_type,
is_inference_mode_enabled,
concat_layout,
accumulate_wgrad_into_main_grad,
)

o = _DownProjection.apply(
Expand All @@ -501,6 +602,7 @@ def moe_general_routing_inputs(
num_activated_expert_per_token_offset,
True, # is_each_token_has_variable_activated_expert
activation_type,
accumulate_wgrad_into_main_grad,
)

return o, expert_frequency
Loading