triton-moe / sanity-layer.py

drbh

feat: update sanity tests and readme

a34e63e over 1 year ago

2.68 kB

# /// script
# dependencies = [
#   "kernels @ git+https://github.com/huggingface/kernels.git@32ec496c5a3169ef02e06a21fb864e0412",
#   "numpy",
#   "torch",
# ]
# ///
import torch
from torch import nn

from kernels import use_kernel_forward_from_hub
from kernels import (
    kernelize,
    register_kernel_mapping,
    use_kernel_mapping,
    LayerRepository,
)

# Make reproducible
torch.manual_seed(42)
torch.cuda.manual_seed(42)


@use_kernel_forward_from_hub("MoE")
class MoE(nn.Module):

    def forward(
        self,
        hidden_states: torch.Tensor,
        router_idx: torch.Tensor,
        router_wt: torch.Tensor,
        alpha: float,
        gate_up_weights: torch.Tensor,
        gate_up_bias: torch.Tensor,
        down_weights: torch.Tensor,
        down_bias: torch.Tensor,
    ) -> torch.Tensor:
        return 0.0


kernel_layer_mapping = {
    "MoE": {
        "cuda": LayerRepository(
            repo_id="kernels-community/triton-moe",
            layer_name="MoE",
            revision="main",
        )
    }
}

layer = MoE()

register_kernel_mapping(kernel_layer_mapping)

with use_kernel_mapping(kernel_layer_mapping):
    layer = kernelize(
        layer,
        needs_backward=True,
        # needs_torch_compile=True,
        use_fallback=False,
        device="cuda",
    )

num_experts = 4
batch_tokens = 8
hidden_size = 16
expert_dim = hidden_size // num_experts
topk = 2


# Example usage
hidden_states = torch.randn(
    batch_tokens,
    hidden_size,
    device="cuda",
    dtype=torch.float32,
    requires_grad=True,
)
router_idx = torch.randint(0, num_experts, (batch_tokens, topk), device="cuda")
router_wt = torch.rand(batch_tokens, topk, device="cuda")
router_wt = router_wt / router_wt.sum(dim=-1, keepdim=True)  # Normalize

# Initialize parameters (shared between both implementations)
gate_up_proj = nn.Parameter(
    torch.randn(num_experts, hidden_size, 2 * expert_dim, device="cuda")
)
gate_up_proj_bias = nn.Parameter(
    torch.randn(num_experts, 2 * expert_dim, device="cuda")
)
down_proj = nn.Parameter(
    torch.randn(num_experts, expert_dim, hidden_size, device="cuda")
)
down_proj_bias = nn.Parameter(torch.randn(num_experts, hidden_size, device="cuda"))
alpha = 1.702


# Run the MoE layer
output = layer(
    hidden_states,
    router_idx,
    router_wt,
    alpha,
    gate_up_proj,
    gate_up_proj_bias,
    down_proj,
    down_proj_bias,
)

print("Output sum:", output.sum().item())
print("Output shape:", output.shape)

loss = output.sum()
loss.backward()
print("Loss:", loss.item())

print("Gate up weights grad shape:", gate_up_proj.grad.shape)
print("Gate up weights grad shape:", gate_up_proj_bias.grad)

	# /// script
	# dependencies = [
	# "kernels @ git+https://github.com/huggingface/kernels.git@32ec496c5a3169ef02e06a21fb864e0412",
	# "numpy",
	# "torch",
	# ]
	# ///
	import torch
	from torch import nn

	from kernels import use_kernel_forward_from_hub
	from kernels import (
	kernelize,
	register_kernel_mapping,
	use_kernel_mapping,
	LayerRepository,
	)

	# Make reproducible
	torch.manual_seed(42)
	torch.cuda.manual_seed(42)


	@use_kernel_forward_from_hub("MoE")
	class MoE(nn.Module):

	def forward(
	self,
	hidden_states: torch.Tensor,
	router_idx: torch.Tensor,
	router_wt: torch.Tensor,
	alpha: float,
	gate_up_weights: torch.Tensor,
	gate_up_bias: torch.Tensor,
	down_weights: torch.Tensor,
	down_bias: torch.Tensor,
	) -> torch.Tensor:
	return 0.0


	kernel_layer_mapping = {
	"MoE": {
	"cuda": LayerRepository(
	repo_id="kernels-community/triton-moe",
	layer_name="MoE",
	revision="main",
	)
	}
	}

	layer = MoE()

	register_kernel_mapping(kernel_layer_mapping)

	with use_kernel_mapping(kernel_layer_mapping):
	layer = kernelize(
	layer,
	needs_backward=True,
	# needs_torch_compile=True,
	use_fallback=False,
	device="cuda",
	)

	num_experts = 4
	batch_tokens = 8
	hidden_size = 16
	expert_dim = hidden_size // num_experts
	topk = 2


	# Example usage
	hidden_states = torch.randn(
	batch_tokens,
	hidden_size,
	device="cuda",
	dtype=torch.float32,
	requires_grad=True,
	)
	router_idx = torch.randint(0, num_experts, (batch_tokens, topk), device="cuda")
	router_wt = torch.rand(batch_tokens, topk, device="cuda")
	router_wt = router_wt / router_wt.sum(dim=-1, keepdim=True) # Normalize

	# Initialize parameters (shared between both implementations)
	gate_up_proj = nn.Parameter(
	torch.randn(num_experts, hidden_size, 2 * expert_dim, device="cuda")
	)
	gate_up_proj_bias = nn.Parameter(
	torch.randn(num_experts, 2 * expert_dim, device="cuda")
	)
	down_proj = nn.Parameter(
	torch.randn(num_experts, expert_dim, hidden_size, device="cuda")
	)
	down_proj_bias = nn.Parameter(torch.randn(num_experts, hidden_size, device="cuda"))
	alpha = 1.702


	# Run the MoE layer
	output = layer(
	hidden_states,
	router_idx,
	router_wt,
	alpha,
	gate_up_proj,
	gate_up_proj_bias,
	down_proj,
	down_proj_bias,
	)

	print("Output sum:", output.sum().item())
	print("Output shape:", output.shape)

	loss = output.sum()
	loss.backward()
	print("Loss:", loss.item())

	print("Gate up weights grad shape:", gate_up_proj.grad.shape)
	print("Gate up weights grad shape:", gate_up_proj_bias.grad)