⚡ Yoshida: Production-Grade General LLM Reinforcement Learning Framework

Yoshida is a high-performance, modular General LLM Reinforcement Learning (RL) and Post-Training Framework built natively in pure PyTorch with Zero External RL Library Lock-in (No TRL, No Ray, No Deepspeed-RL dependencies).

Designed for both agile single-GPU research and multi-GPU cluster scaling, Yoshida supports DeepSeek-R1 style GRPO (Reasoning & Math RL), Online PPO (Actor-Critic), Direct Preference Optimization (DPO), Bradley-Terry Reward Modeling (RM), and Supervised Fine-Tuning (SFT).

It pairs exact mathematical formulations with modern spectral optimizers (Dion3 with Gram Newton-Schulz, Muon, Fused AdamW), colocated in-process vLLM PagedAttention rollouts (3,300+ tokens/sec), native multi-GPU FSDP (ZeRO-3), and local structured experiment workspaces.

---

📚 Documentation

Explore our comprehensive guides in docs/:

Guide	Description
🚀 Getting Started Guide	Step-by-step installation, hardware verification, and your first 30-second training run.
📊 Dataset Preparation	Standard JSON/JSONL formatting rules for SFT, DPO, RM, and GRPO reasoning/alignment datasets.
⚙️ Configuration Guide	Declarative PyTorch Lightning-style YAML configuration and dynamic dot-notation CLI overrides.
🍳 Training Recipes	Ready-to-use copy-paste recipes for GRPO reasoning RL, DPO, Multi-GPU FSDP, and vLLM acceleration.
🏛️ SLURM Cluster & HPC Guide	Multi-node torchrun batch scripts, InfiniBand tuning, and spot preemption auto-resuming.
🛡️ Evaluation & Benchmarks	Automated multi-axis evaluation for reasoning accuracy, instruction following, and guardrails.
❓ Troubleshooting & FAQ	Resolving CUDA Out-of-Memory (OOM), optimizer selection guidance, and NCCL tuning.

---

⚡ Key Architectural Features

• Zero RL Library Lock-in: Exact mathematical formulations of:
  • GRPO: DeepSeek R1-style Group Relative Policy Optimization with Schulman $k_3$ non-negative KL penalty.
  • Online PPO: Generalized Advantage Estimation ($\text{GAE}(\gamma, \lambda)$) with Value Head Critic.
  • DPO: Direct Preference Optimization with zero-memory reference switching.
  • Bradley-Terry Reward Modeling: Pairwise log-sigmoid preference scoring.
  • SFT: Masked sequence cross-entropy.
• High-Throughput Rollout Subsystem (align_core/rollout/):
  • Colocated vLLM Engine: Continuous batching and PagedAttention delivering 3,300+ tokens/second in-process on a single GPU.
  • Native PyTorch Engine: Zero-dependency FlashAttention SDPA baseline (190 TPS) that runs anywhere without serving daemons.
  • First-Class Trajectory Provenance: Full tracking of behavior policy $\mu(a|s)$, policy_version, base_id, and exact importance ratios.
• Modern Spectral Optimizers:
  • Dion3: Gram Newton-Schulz (GNS) with Fractional Row Selection (up to $6\times$ lighter than vanilla Muon with identical validation quality).
  • Muon: Quintic Newton-Schulz matrix polar decomposition (NeMo-RL aligned).
  • Fused AdamW: PyTorch native single-kernel CUDA acceleration.
• Multi-GPU Distributed Engine: Native torchrun support with DDP and FSDP (ZeRO-3) parameter + gradient + optimizer state sharding with auto-wrapping transformer decoder layers.
• Local-First Experiment Workspace: Structured <project_slug>/<experiment_slug> folders containing config.yaml, metrics.jsonl, summary.json, rotating checkpoints, and synchronized WandB namespaces.
• Fault-Tolerant Checkpointing: Exact state restoration (weights, AdamW/Muon moments, LR scheduler curves, AMP scalers, CPU/CUDA RNG seeds) with graceful SIGINT/SIGTERM preemption emergency saving.
• Hardware Profiler & Smoketest: Empirical GEMM matrix benchmarks (TOPS / TFLOPs), tokens/sec, and Model FLOPs Utilization (MFU %).

---

🚀 Quickstart

1. Hardware Verification

python smoketest.py --config configs/smoketest.yaml

2. Run DeepSeek-Style GRPO Reasoning RL (with Colocated vLLM)

python train.py \
  --method grpo \
  --rollout_backend vllm \
  --optimizer dion3 \
  --project llm-reasoning \
  --experiment grpo-qwen-reasoning \
  --group_size 4

3. Run Direct Preference Optimization (DPO) with Dion3

python train.py \
  --method dpo \
  --optimizer dion3 \
  --project llm-alignment \
  --experiment dpo-dion3-pilot \
  --max_steps 20

4. Multi-GPU Distributed Training (FSDP ZeRO-3)

torchrun --nproc_per_node=4 train.py --config configs/fsdp_grpo_safety.yaml

5. Benchmark & Evaluation

python evaluate.py \
  --base_model Qwen/Qwen3-0.6B \
  --sft_adapter ./experiments/llm-alignment/sft-baseline/final_model \
  --dpo_adapter ./experiments/llm-alignment/dpo-dion3-pilot/final_model \
  --grpo_adapter ./experiments/llm-reasoning/grpo-qwen-reasoning/final_model

---

🐳 Docker & Container Deployment

# Build Docker image
docker build -t yoshida:latest .

# Run hardware smoketest & TFLOPs profiler
docker compose run --rm smoketest

# Run GRPO RL Training with GPU passthrough
docker compose run --rm train-grpo

# Run Multi-GPU FSDP (ZeRO-3) across all GPUs
docker compose run --rm train-fsdp

---

📂 Repository Structure

yoshida/
├── align_core/             # Core Framework Library
│   ├── rollout/            # RolloutManager, PolicyVersion, Trajectory provenance, and backends (vLLM, NativeTorch)
│   ├── configs/            # Dataclass & PyTorch Lightning-style YAML config parser
│   ├── distributed/        # Multi-GPU context, NCCL collective comms, DDP & FSDP (ZeRO-3)
│   ├── hardware/           # NVIDIA GPU auto-detection, TF32, SDPA, FLOPs/MFU profiler
│   ├── logging/            # Workspace slug manager, local JSONL streaming, & WandB tracker
│   ├── models/             # Native PyTorch loss math (GRPO, DPO, PPO/GAE, BT, SFT)
│   ├── trainers/           # SFT, RM, DPO, PPO, GRPO trainers, CheckpointManager & Optimizers (Dion3/Muon)
│   └── data/               # Safety datasets, reasoning prompts & collators
├── configs/                # Ready-to-use YAML configuration files
│   ├── dpo_safety.yaml     # DPO configuration preset
│   ├── grpo_safety.yaml    # GRPO configuration preset
│   ├── fsdp_grpo_safety.yaml # Multi-GPU FSDP (ZeRO-3) preset
│   └── smoketest.yaml      # Hardware profiling preset
├── docs/                   # Documentation Suite
│   ├── GETTING_STARTED.md
│   ├── DATASET_PREPARATION.md
│   ├── CONFIGURATION_GUIDE.md
│   ├── TRAINING_RECIPES.md
│   ├── EVALUATION_AND_BENCHMARKS.md
│   └── TROUBLESHOOTING_AND_FAQ.md
├── train.py                # Unified CLI & torchrun trainer
├── evaluate.py             # Evaluation benchmark
├── smoketest.py            # Hardware GEMM TOPS/TFLOPs & training simulation profiler
├── Dockerfile              # Production NVIDIA CUDA / PyTorch image
└── docker-compose.yml      # Multi-GPU container orchestration

---

📄 License

MIT License