jhhj25/reef-nccl-ep-qwen3-8b-lora-discovery-16step

reef NCCL EP discovery — Qwen3-8B LoRA, 16-step RL run (step 8 checkpoint)

This repository holds the surviving checkpoint from a 16-step reinforcement-learning training run in which a Qwen3-8B model learned to propose configurations for an expert-parallel collective-communication kernel (dispatch + combine). The reward for each proposed configuration was its bandwidth measured on real hardware (two nodes, 8x H200 GPUs each), gated on a 16/16-pass correctness test; configurations that failed correctness received no reward.

What is in this repo

Only the step 8 adapter survived a storage incident: the training work directory was deleted by a follow-on job before the checkpoints were archived. The step-8 adapter was recovered from the memory of a still-running inference server that had the file mapped, and validates as a complete PEFT LoRA safetensors file (504 tensors, all 36 layers, shapes consistent with Qwen3-8B). The other 15 checkpoints are lost. The full training curve is reproduced below from the run logs.

• —step_8/ — PEFT LoRA adapter, rank 32, alpha 32, dropout 0, target modules qproj, kproj, vproj, oproj, gateproj, upproj, down_proj.

The adapter_config.json was reconstructed from the training recipe configuration (the original file was lost with the directory); the safetensors weights are byte-identical across two independent memory dumps (sha256 93d717a6d120d1a12faa5dbbfa34eb47e34e1b77c97bc846471b8246f89da713).

Training recipe

• —Base model: Qwen/Qwen3-8B, PEFT LoRA rank 32.
• —Trainer: reef TTTD recipe (test-time training / discovery) on a slime + Megatron stack, grouped relative-reward RL: 4 groups x 4 rollouts per step, 16 steps.
• —Task: emit an integer-valued kernel configuration (queue depth, chunk size, split factor, combine block, SM count) for an expert-parallel dispatch+combine collective. Reward = measured dispatch + combine bandwidth (GB/s total) on hardware, with a hard correctness gate.

Reward curve (per-step mean of rollout rewards, GB/s)

Best single episode: 175.69 GB/s (config q2 c256 sp32 cb64 sms12, correctness 16/16). For scale: a 480B-parameter model prompted zero-shot on the same task and scored by the same judge reached 147.46 GB/s, so the best configuration discovered by the trained 8B model exceeds the zero-shot proposal of the much larger model.

Loading

from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

base = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-8B", torch_dtype="bfloat16")
model = PeftModel.from_pretrained(
    base,
    "jhhj25/reef-nccl-ep-qwen3-8b-lora-discovery-16step",
    subfolder="step_8",
)
tok = AutoTokenizer.from_pretrained("Qwen/Qwen3-8B")