replicate/megablocks

1# /// script2# requires-python = "==3.10"3# dependencies = [4#     "numpy",5#     "kernels",6#     "torch"7# ]8# ///9 10import torch11from collections import namedtuple12 13from kernels import get_kernel14 15# Make reproducible16torch.manual_seed(42)17torch.cuda.manual_seed(42)18 19# Download optimized kernels from the Hugging Face hub20megablocks = get_kernel("kernels-community/megablocks")21print("MegaBlocks kernel downloaded successfully.")22 23model = megablocks.layers.MegaBlocksMoeMLP()24model.experts = namedtuple("Experts", ["gate_up_proj", "gate_down_proj", "down_proj", "hidden_size"])25print("MegaBlocksMoeMLP instance created successfully.")26 27# Config28ne, hs, isz = 128, 1152, 307229 30# Router with proper initialization31model.router = torch.nn.Linear(hs, ne, device="cuda")32torch.nn.init.kaiming_uniform_(model.router.weight)33 34# Expert layers with realistic weights35e = model.experts36e.gate_up_proj = torch.nn.Parameter(torch.randn(ne, hs, isz, device="cuda") * 0.02)37e.gate_up_proj_bias = torch.nn.Parameter(torch.zeros(ne, isz, device="cuda"))38e.down_proj = torch.nn.Parameter(torch.randn(ne, 1536, hs, device="cuda") * 0.02)39e.down_proj_bias = torch.nn.Parameter(torch.zeros(ne, hs, device="cuda"))40e.hidden_size = hs41print("Expert layers initialized successfully.")42 43# Test with normalized input44x = torch.randn(1, 1, hs, device="cuda") * 0.145output, expert_weights = model(x)46print("Model forward pass completed successfully.")47 48print(f"Output shape: {output.shape}")49print(f"Output range: [{output.min():.3f}, {output.max():.3f}]")50print(f"Output: {output.flatten()[:10]}")51print(f"Expert weights sum: {expert_weights.sum():.3f}")