290 lines
10 KiB
Markdown
290 lines
10 KiB
Markdown
# RADLADS in Foundry
|
|
|
|
> Adding transformer-to-RWKV distillation tooling to Foundry as a first-class tool.
|
|
|
|
---
|
|
|
|
## What is RADLADS?
|
|
|
|
[RADLADS](https://arxiv.org/abs/2505.03005) (Rapid Attention Distillation to Linear Attention Decoders at Scale) is a protocol for converting pretrained transformer models into RWKV linear attention models. Developed by [Recursal](https://github.com/recursal/RADLADS-paper) (Apache 2.0 license).
|
|
|
|
It has been proven at scale:
|
|
- Converted Qwen2.5 models (7B, 32B, 72B) into RWKV-6 and RWKV-7 variants
|
|
- Uses only 350-700M tokens (~0.005% of original pretraining data)
|
|
- Retains ~95-98% of teacher model quality
|
|
- The 72B conversion cost less than $2,000
|
|
|
|
Currently implemented in Python/PyTorch. We want to bring this capability into Foundry as pure C.
|
|
|
|
---
|
|
|
|
## Why Add RADLADS to Foundry?
|
|
|
|
### 1. CPU Cluster Training
|
|
|
|
The original RADLADS code assumes GPU (A100/H100). A Foundry implementation enables:
|
|
- Training on CPU clusters with 8GB+ per node
|
|
- Using our existing Docker Swarm infrastructure
|
|
- Model parallelism via Foundry's `distributed/` and `swarm/` modules
|
|
- Quantized teacher inference (Q8) to fit in memory
|
|
|
|
### 2. Zero External Dependencies
|
|
|
|
Foundry + JACLibc means no Python, no PyTorch, no CUDA toolkit required. The entire conversion pipeline runs as a single compiled binary or across a distributed cluster.
|
|
|
|
### 3. Tighter Memory Control
|
|
|
|
C gives us byte-level control over memory layout. Critical for fitting a 7B teacher + student in 8GB per node via:
|
|
- Memory-mapped weights (safetensors mmap)
|
|
- Layer-by-layer processing (don't load full model)
|
|
- Mixed precision (Q8 frozen weights, FP32 trainable)
|
|
- Gradient checkpointing
|
|
|
|
### 4. BootAI / Bare Metal Integration
|
|
|
|
A Foundry-native RADLADS means the conversion tooling can eventually run on BootAI itself — converting models directly on bare metal hardware.
|
|
|
|
---
|
|
|
|
## What RADLADS Needs from Foundry
|
|
|
|
### Already Implemented
|
|
|
|
| Foundry Module | RADLADS Use | Status |
|
|
|----------------|-------------|--------|
|
|
| `tensor/` | All tensor ops, autograd, backward pass | ✅ |
|
|
| `attention/` | Multi-head attention, KV cache, RoPE | ✅ |
|
|
| `safetensors/` | Load teacher/student weights, save checkpoints | ✅ |
|
|
| `tokenizer/` | BPE tokenizer (load Qwen HuggingFace JSON) | ✅ |
|
|
| `data/` | JSONL data loading, batching, shuffling | ✅ |
|
|
| `optim/` | Adam/AdamW optimizers, cosine scheduler | ✅ |
|
|
| `distributed/` | All-reduce, gradient accumulation | ✅ |
|
|
| `swarm/` | Parameter server, gradient compression, Docker Swarm | ✅ |
|
|
| `quantization/` | INT8/INT4 for frozen teacher weights | ✅ |
|
|
| `progress/` | Training progress bars, metrics | ✅ |
|
|
| `lora/` | LoRA layers (may be useful for parameter-efficient stages) | ✅ |
|
|
| `eval/` | Perplexity evaluation between runs | ✅ |
|
|
|
|
### Not Yet Implemented
|
|
|
|
| Need | Description | Where |
|
|
|------|-------------|-------|
|
|
| **Transformer forward pass** | Qwen2.5 architecture (RMSNorm, SwiGLU, GQA, RoPE) | `transformer/` module (planned) |
|
|
| **RWKV-7 forward pass** | WKV-7 kernel, token shift, decay/key/output gates | `transformer/` module (planned) |
|
|
| **RWKV-X sparse attention** | Top-k chunk sparse attention with bounded KV cache | `transformer/` module |
|
|
| **KL divergence loss** | Already have `kl_div` in tensor — verify it supports distillation use | `tensor/` |
|
|
| **Pipeline parallelism** | Split model layers across nodes, pass activations | `distributed/` or `swarm/` |
|
|
| **Mixed-precision training** | FP32 gradients on Q8 frozen weights | `tensor/` + `quantization/` |
|
|
|
|
---
|
|
|
|
## Proposed Module: `radlads/`
|
|
|
|
A new top-level Foundry module implementing the three-step RADLADS pipeline.
|
|
|
|
### Structure
|
|
|
|
```
|
|
foundry/
|
|
└── radlads/
|
|
├── radlads.h # Public API
|
|
├── radlads.c # Pipeline orchestration
|
|
├── alignment.c # Step 1: Hidden state alignment
|
|
├── distillation.c # Step 2: KL divergence distillation
|
|
├── extension.c # Step 3: Context length extension
|
|
├── weight_mapping.c # Transformer -> RWKV weight transfer
|
|
├── config.h # Conversion configuration
|
|
└── tests/
|
|
└── test_radlads.c
|
|
```
|
|
|
|
### Public API (Draft)
|
|
|
|
```c
|
|
#include "radlads/radlads.h"
|
|
|
|
// Configuration
|
|
typedef struct {
|
|
const char* teacher_path; // Path to teacher safetensors
|
|
const char* student_output; // Path for student output
|
|
const char* tokenizer_path; // Path to tokenizer.json
|
|
const char* data_path; // Path to training data
|
|
|
|
// Architecture
|
|
fr_arch_t teacher_arch; // FR_ARCH_QWEN2, FR_ARCH_LLAMA, etc.
|
|
fr_arch_t student_arch; // FR_ARCH_RWKV7, FR_ARCH_RWKVX
|
|
|
|
// Training
|
|
int alignment_tokens; // Step 1 token budget (default: 100M)
|
|
int distillation_tokens; // Step 2 token budget (default: 500M)
|
|
int extension_tokens; // Step 3 token budget (default: 100M)
|
|
int extension_context_len; // Target context length (default: 16384)
|
|
int batch_size; // Per-node batch size
|
|
float learning_rate; // Default: 3e-4
|
|
|
|
// Memory
|
|
bool quantize_teacher; // Use Q8 for teacher (saves ~50% RAM)
|
|
bool gradient_checkpointing; // Trade compute for memory
|
|
|
|
// Distributed
|
|
int world_size; // Number of nodes (0 = auto-detect)
|
|
fr_parallel_t parallel_mode; // FR_DATA_PARALLEL, FR_PIPELINE_PARALLEL
|
|
} fr_radlads_config;
|
|
|
|
// Run full conversion pipeline
|
|
int fr_radlads_convert(fr_radlads_config* config);
|
|
|
|
// Run individual steps
|
|
int fr_radlads_align(fr_radlads_config* config); // Step 1 only
|
|
int fr_radlads_distill(fr_radlads_config* config); // Step 2 only
|
|
int fr_radlads_extend(fr_radlads_config* config); // Step 3 only
|
|
|
|
// Weight mapping (no training, just transfer weights)
|
|
int fr_radlads_init_student(
|
|
const char* teacher_path,
|
|
const char* student_path,
|
|
fr_arch_t teacher_arch,
|
|
fr_arch_t student_arch
|
|
);
|
|
```
|
|
|
|
### Usage Example
|
|
|
|
```c
|
|
#include "radlads/radlads.h"
|
|
|
|
int main(void) {
|
|
fr_radlads_config config = {
|
|
.teacher_path = "models/Qwen2.5-7B-Instruct",
|
|
.student_output = "models/QRWKV7-7B",
|
|
.tokenizer_path = "models/Qwen2.5-7B-Instruct/tokenizer.json",
|
|
.data_path = "data/minipile.jsonl",
|
|
.teacher_arch = FR_ARCH_QWEN2,
|
|
.student_arch = FR_ARCH_RWKV7,
|
|
.alignment_tokens = 100000000,
|
|
.distillation_tokens = 500000000,
|
|
.extension_tokens = 100000000,
|
|
.extension_context_len = 16384,
|
|
.batch_size = 1,
|
|
.learning_rate = 3e-4,
|
|
.quantize_teacher = true,
|
|
.gradient_checkpointing = true,
|
|
.world_size = 0, // auto-detect
|
|
.parallel_mode = FR_PIPELINE_PARALLEL
|
|
};
|
|
|
|
return fr_radlads_convert(&config);
|
|
}
|
|
```
|
|
|
|
---
|
|
|
|
## Dependencies on Transformer Module
|
|
|
|
RADLADS requires model forward passes for both teacher and student architectures. This means the `transformer/` module (currently missing) must be built first. The needed architectures:
|
|
|
|
### Teacher Architectures (inference only, no gradients)
|
|
|
|
| Architecture | Components | Priority |
|
|
|-------------|------------|----------|
|
|
| Qwen2.5 | RMSNorm, SwiGLU FFN, GQA, RoPE | High (our target teacher) |
|
|
| Llama 3 | RMSNorm, SwiGLU FFN, GQA, RoPE | Medium (nearly identical to Qwen) |
|
|
|
|
### Student Architectures (training, needs gradients)
|
|
|
|
| Architecture | Components | Priority |
|
|
|-------------|------------|----------|
|
|
| RWKV-7 | WKV-7 kernel, token shift, decay/key/output gates, GroupNorm, ReLU-squared FFN | High |
|
|
| RWKV-X | RWKV-7 blocks + Top-k chunk sparse attention + bounded KV cache | High |
|
|
|
|
### Build Order
|
|
|
|
```
|
|
transformer/ (core building blocks)
|
|
├── rmsnorm, layernorm, embedding
|
|
├── swiglu_ffn, relu_squared_ffn
|
|
├── gqa_attention (+ RoPE)
|
|
├── wkv7_kernel
|
|
├── rwkv7_block (token shift, decay, gates)
|
|
└── sparse_attention_block (top-k chunk, bounded cache)
|
|
│
|
|
▼
|
|
radlads/ (conversion pipeline)
|
|
├── weight_mapping (transformer -> RWKV weight init)
|
|
├── alignment (step 1)
|
|
├── distillation (step 2)
|
|
└── extension (step 3)
|
|
```
|
|
|
|
---
|
|
|
|
## Memory Budget for CPU Cluster
|
|
|
|
Targeting 8GB-per-node machines:
|
|
|
|
### Option A: Pipeline Parallel (3 nodes)
|
|
|
|
```
|
|
Node 0: Teacher layers 0-9 (Q8) + Student layers 0-9 ~7 GB
|
|
Node 1: Teacher layers 10-19 (Q8) + Student layers 10-19 ~7 GB
|
|
Node 2: Teacher layers 20-27 (Q8) + Student layers 20-27 ~5 GB
|
|
+ Embeddings + LM head + Optimizer state
|
|
```
|
|
|
|
Activations passed between nodes via TCP (swarm module).
|
|
|
|
### Option B: Layer-by-Layer Alignment (1 node at a time)
|
|
|
|
For Step 1 (alignment), we only need one teacher layer and one student layer at a time:
|
|
```
|
|
Load teacher layer N (Q8): ~250 MB
|
|
Load student layer N (FP32): ~500 MB
|
|
Activations + gradients: ~200 MB
|
|
Optimizer state: ~1 GB
|
|
Total per layer: ~2 GB (fits easily in 8 GB)
|
|
```
|
|
|
|
Process layers sequentially. This is the most memory-efficient approach for Step 1.
|
|
|
|
### Option C: Hybrid
|
|
|
|
- Step 1 (alignment): Layer-by-layer on single node (Option B)
|
|
- Step 2 (distillation): Pipeline parallel across 3+ nodes (Option A)
|
|
- Step 3 (extension): Same as Step 2
|
|
|
|
---
|
|
|
|
## Estimated Timeline
|
|
|
|
| Phase | Work | Depends On |
|
|
|-------|------|------------|
|
|
| 1. `transformer/` core blocks | RMSNorm, SwiGLU, GQA, embedding | tensor/ (done) |
|
|
| 2. `transformer/` Qwen2.5 | Full Qwen forward pass (inference) | Phase 1 |
|
|
| 3. `transformer/` RWKV-7 | WKV-7 kernel, full block (with autograd) | Phase 1 |
|
|
| 4. `transformer/` RWKV-X | Block expansion, sparse attention | Phase 3 |
|
|
| 5. `radlads/` weight mapping | Init student from teacher weights | Phase 2, 3 |
|
|
| 6. `radlads/` alignment | Step 1 hidden state matching | Phase 5 |
|
|
| 7. `radlads/` distillation | Step 2 KL divergence | Phase 6 |
|
|
| 8. `radlads/` extension | Step 3 long context | Phase 7 |
|
|
| 9. Pipeline parallelism | Multi-node model splitting | distributed/ (done) |
|
|
| 10. Integration testing | End-to-end Qwen -> RWKV-X | All above |
|
|
|
|
---
|
|
|
|
## References
|
|
|
|
- [RADLADS Paper](https://arxiv.org/abs/2505.03005) (May 2025)
|
|
- [RADLADS Code](https://github.com/recursal/RADLADS-paper) (Apache 2.0)
|
|
- [RADLADS Models](https://huggingface.co/collections/recursal/radlads-6818ee69e99e729ba8a87102)
|
|
- [RWKV-X Paper](https://arxiv.org/abs/2504.21463) (April 2025)
|
|
- [HALO: Hybrid Linear Attention Done Right](https://arxiv.org/abs/2601.22156) (January 2026)
|
|
- [MiniCPM-SALA](https://arxiv.org/abs/2602.11761) (February 2026)
|
|
|
|
---
|
|
|
|
## Changelog
|
|
|
|
| Date | Change |
|
|
|------|--------|
|
|
| 2026-02-19 | Initial proposal |
|