Bumblebee

A lightweight distributed LLM training library with 5D parallelism (TP × EP × ETP × CP × PP), ZeRO-2/3 optimizer offload, and a protocol-driven model registry.

Architecture

Bumblebee has three layers, enforced by AST-level import boundaries:

┌─────────────────────────────────────────────────┐
│  runtime/          # Session, backends, config   │
│  ├─ backends/bb/   # Megatron5D (default)        │
│  └─ backends/bridge/ # mbridge hybrid            │
├─────────────────────────────────────────────────┤
│  primitive/        # Framework-agnostic building  │
│  ├─ parallel/      # 5D state, linear, pipeline   │
│  ├─ modules/       # GQA, MLA, MoE, DenseMLP      │
│  ├─ optimizers/    # ZeRO-2/3 + CPU offload       │
│  ├─ recompute.py   # Activation checkpointing     │
│  ├─ scheduler.py   # Cosine/Warmup/WSD LR         │
│  ├─ device.py      # Device abstraction           │
│  └─ train_step.py  # Fwd/bwd/grad-norm/clip       │
├─────────────────────────────────────────────────┤
│  model/            # Model-specific composition    │
│  ├─ qwen3_moe/     # Qwen3-MoE (lite/hybrid/bridge)│
│  ├─ deepseek_v3/   # DeepSeek-V3 (MLA + shared exp)│
│  └─ registry.py    # Protocol registration        │
└─────────────────────────────────────────────────┘

Key principle: primitive/ never imports model/ or runtime/. model/ imports primitive/ only. runtime/ imports both. This is enforced by tests/test_arch_boundary.py.

Quick Start

Basic Training (BB Backend)

from bumblebee import BBConfig, ParallelConfig, OptimizerConfig

cfg = BBConfig(
    model_name="qwen3_moe",          # or "deepseek_v3"
    impl="lite",                     # "lite" | "hybrid" | "bridge"
    hf_path="/path/to/hf/weights",
    parallel=ParallelConfig(tp=2, ep=4, pp=2, cp=1),
    optimizer=OptimizerConfig(
        lr=3e-4,
        lr_decay_style="cosine",
        lr_warmup_steps_ratio=0.05,
        total_training_steps=10000,
        offload_fraction=0.3,        # ZeRO-2 CPU offload
    ),
)

from bumblebee.runtime.backends.bb.runtime import Megatron5DRuntime
rt = Megatron5DRuntime(hf_path=cfg.hf_path, cfg=cfg)
handle = rt.build_model()

# Training loop
from bumblebee.primitive.data import infinite_batches
from bumblebee.primitive.scheduler import build_scheduler

sched = build_scheduler("cosine_warmup", handle.optimizer,
                        base_lr=3e-4, total_steps=10000,
                        warmup_steps=500, min_lr=3e-5)

for step, batch in enumerate(infinite_batches(vocab_size=151936, seq_len=4096)):
    rt.forward_backward(handle, batch, num_microbatches=4)
    rt.optimizer_step(handle)
    sched.step()
    if step >= 10000:
        break

Adding a New Model

1. Create bumblebee/model/<name>/config.py — a @dataclass with architecture params.
2. Create bumblebee/model/<name>/lite/protocol.py — implement:
   • ImplConfig (per-impl knobs)
   • build_model_config(source, **overrides) → ModelConfig
   • build_model(model_cfg, *, impl_cfg) → ModelBundle
   • Optional: vocab_size(), activated_params(), load_hf_weights()
3. Create bumblebee/model/<name>/model.py — the nn.Module using primitive/ building blocks.
4. Register in bumblebee/model/registry.py:

   register_model("my_model", package="bumblebee.model.my_model",
       hf_model_types=["my_model_type"],
       impls={"lite": "bumblebee.model.my_model.lite.protocol"})

See bumblebee/model/deepseek_v3/ for a reference implementation using MLA + shared experts.

Parallelism

Bumblebee supports 5D parallelism via a dual-rank decomposition:

Dim	Description	Process Group
TP	Tensor Parallel	tp_group
EP	Expert Parallel	ep_group
ETP	Expert Tensor Parallel	etp_group
CP	Context Parallel	cp_group
PP	Pipeline Parallel	pp_group

The dual-rank decomposition separates dense and expert parameter groups:

• Dense params: sharded across dp_group = world / (pp * cp)
• Expert params: sharded across ep_dp_group = world / (pp * cp * ep)

ETP (expert tensor parallel) further shards expert weights within a node, enabling large expert counts without OOM.

Optimizer

DistributedOptimizer implements ZeRO-2 with optional CPU offload:

• Gradient ReduceScatter: replaces AllReduce, shards gradients across DP group
• CPU Offload (offload_fraction): moves master weights + Adam state to pinned CPU
• Gradient Offload (grad_offload_fraction): D2H gradient copies to save GPU memory
• Async H2D: uses CUDA streams with non_blocking=True to overlap copy with compute
• Fused Adam: torch._foreach_* for batched parameter updates

from bumblebee.primitive.optimizers.native import DistributedOptimizer

opt = DistributedOptimizer(
    model_params,
    ps=parallel_state,
    lr=3e-4,
    offload_fraction=0.3,       # 30% of optimizer state on CPU
    grad_offload_fraction=0.2,  # 20% of gradients offloaded
)

LR Schedulers

from bumblebee.primitive.scheduler import build_scheduler

# Cosine with linear warmup
sched = build_scheduler("cosine_warmup", optimizer,
    base_lr=3e-4, total_steps=10000,
    warmup_steps=500, min_lr=3e-5)

# WSD (Warmup-Stable-Decay) — DeepSeek-V3 style
sched = build_scheduler("wsd", optimizer,
    base_lr=3e-4, total_steps=10000,
    warmup_steps=200, stable_steps=8000, min_lr=0.0,
    decay_type="cosine")

Supported: constant, cosine, cosine_warmup, wsd.

Device Abstraction

Bumblebee uses get_device() instead of hardcoding "cuda", enabling CPU-only testing:

from bumblebee.primitive.device import get_device

# On GPU machines: "cuda", on CPU-only: "cpu"
device = get_device()
t = torch.empty(shape, device=device)

Set BUMBLEBEE_FORCE_CPU=1 to force CPU even on GPU machines.

CPU-Safe Dependency Shims

transformer_engine and megatron.core are lazy-loaded with CPU-safe fallbacks:

• On GPU machines: real TE/MC modules are used (full FP8, fused kernels, etc.)
• On CPU-only machines: pure-PyTorch shims provide the same API surface (RMSNorm, Linear, GroupedLinear, DotProductAttention, RoPE, etc.) so that protocol modules import without the GPU stack installed

from bumblebee.primitive._te_compat import get_te, has_te
from bumblebee.primitive._mc_compat import get_mc, has_mc

te = get_te()  # real TE or shim
mc = get_mc()  # real MC or shim

Testing

# Run all tests (CPU-only, no GPU/TE/MC required)
python -m pytest tests/ -q

# Run specific suite
python -m pytest tests/runtime/ -q
python -m pytest tests/qwen3_moe/ -q

Project Structure

bumblebee/
├── __init__.py              # Lazy top-level exports
├── runtime/
│   ├── contracts/           # Shared config types
│   └── backends/
│       ├── bb/              # Megatron5D (default)
│       └── bridge/          # mbridge hybrid
├── primitive/
│   ├── parallel/            # 5D state, linear, pipeline, SP
│   ├── modules/             # GQA, MLA, MoE, DenseMLP, Router
│   ├── optimizers/          # ZeRO-2/3 + CPU offload
│   ├── ckpt/                # DCP distributed checkpointing
│   ├── _te_compat.py        # TE lazy loader + CPU shim
│   ├── _mc_compat.py        # MC lazy loader + CPU shim
│   ├── device.py            # Device abstraction
│   ├── scheduler.py         # LR schedulers
│   ├── recompute.py         # Activation checkpointing
│   ├── train_step.py        # Fwd/bwd/grad-norm/clip
│   ├── data.py              # Batch generation
│   ├── bundle.py            # ModelBundle return type
│   ├── protocols.py         # Train-side protocols
│   └── config.py            # HF config loading
├── model/
│   ├── registry.py          # Model registration
│   ├── qwen3_moe/           # Qwen3-MoE model
│   └── deepseek_v3/         # DeepSeek-V3 model
└── benchmarks/              # Throughput benchmarking

Design Principles

1. Protocol-driven: Models implement a minimal protocol (ImplConfig, build_model_config, build_model). The runtime never hardcodes model specifics.
2. AST-enforced boundaries: Import rules are checked by tests, not conventions.
3. CPU-testable: Every module imports and basic tests pass without GPU/TE/MC installed.
4. Composable parallelism: 5D parallelism is a first-class concept, not a post-hoc wrapper.
5. ZeRO-native: The optimizer IS the distributed wrapper — no separate DDP layer.