entrenar/finetune/instruct_pipeline/
mod.rs

1//! Instruction-following fine-tuning pipeline (GH-371)
2//!
3//! Wires Transformer + LoRA for causal language model fine-tuning on
4//! instruction-response pairs.
5//!
6//! # Architecture
7//!
8//! ```text
9//! [prompt_ids ++ response_ids] -> Transformer.forward() -> logits [seq_len, vocab_size]
10//!   -> causal_lm_loss(logits[prompt_len..], response_ids) -> scalar loss
11//! ```
12//!
13//! # Contract
14//!
15//! - F-INST-002: Loss computed only on response tokens (prompt tokens masked)
16//! - F-INST-003: Perplexity = exp(avg_loss) reported per epoch
17//! - F-INST-004: LoRA adapters saved in APR format
18
19mod accessors;
20mod backward;
21mod constructors;
22mod cuda_forward;
23mod cuda_init;
24mod generate;
25mod training;
26mod wgpu;
27
28#[cfg(all(test, feature = "cuda"))]
29mod eval_sync_probe;
30#[cfg(all(test, feature = "cuda"))]
31mod parity_probe;
32#[cfg(test)]
33mod tests;
34#[cfg(test)]
35mod tests_cov3;
36#[cfg(test)]
37mod tests_cov3b;
38
39use crate::lora::LoRALayer;
40use crate::optim::{clip_grad_norm_refs, AdamW, Optimizer};
41use crate::tokenizer::HfTokenizer;
42use crate::train::transformer_trainer::step_profiler::StepProfiler;
43use crate::transformer::{Transformer, TransformerConfig};
44use crate::Tensor;
45use std::path::{Path, PathBuf};
46
47#[cfg(feature = "cuda")]
48use crate::autograd::cuda_training::CudaTrainer;
49#[cfg(feature = "cuda")]
50use crate::gpu::guard::VramGuard;
51#[cfg(feature = "cuda")]
52use crate::transformer::{
53    CudaBlock, CudaBlockScratch, CudaLoraGradWorkspace, GpuLoraOptimizerState,
54};
55#[cfg(feature = "cuda")]
56use trueno_gpu::driver::GpuBuffer;
57
58/// Configuration for instruction fine-tuning.
59#[derive(Debug, Clone)]
60pub struct InstructConfig {
61    /// LoRA rank
62    pub lora_rank: usize,
63    /// LoRA alpha
64    pub lora_alpha: f32,
65    /// Learning rate
66    pub learning_rate: f32,
67    /// Number of training epochs
68    pub epochs: usize,
69    /// Maximum sequence length (prompt + response)
70    pub max_seq_len: usize,
71    /// Maximum gradient norm for clipping
72    pub gradient_clip_norm: Option<f32>,
73    /// Quantize frozen weights to NF4 (4-bit) for QLoRA training (default: false).
74    ///
75    /// When enabled, uses `CudaNf4TransformerBlock` (~8x VRAM compression) instead
76    /// of `CudaTransformerBlock`. GPU backward pass updates only LoRA adapters.
77    pub quantize_nf4: bool,
78}
79
80impl Default for InstructConfig {
81    fn default() -> Self {
82        Self {
83            lora_rank: 16,
84            lora_alpha: 32.0,
85            learning_rate: 2e-4,
86            epochs: 3,
87            max_seq_len: 512,
88            gradient_clip_norm: Some(1.0),
89            quantize_nf4: false,
90        }
91    }
92}
93
94/// Result of processing one instruction-response pair.
95#[derive(Debug, Clone)]
96pub struct InstructStepResult {
97    /// Cross-entropy loss on response tokens
98    pub loss: f32,
99    /// Number of response tokens
100    pub num_response_tokens: usize,
101    /// Perplexity = exp(loss)
102    pub perplexity: f32,
103}
104
105/// Result of processing a mini-batch of instruction samples.
106#[derive(Debug, Clone)]
107pub struct InstructBatchResult {
108    /// Average cross-entropy loss across the batch (response tokens only)
109    pub avg_loss: f32,
110    /// Total response tokens in batch
111    pub total_response_tokens: usize,
112    /// Perplexity = exp(avg_loss)
113    pub perplexity: f32,
114    /// Gradient norm before clipping
115    pub grad_norm: f32,
116}
117
118/// Instruction fine-tuning pipeline.
119///
120/// Owns the transformer and LoRA adapters. Uses `Transformer::forward()`
121/// for causal LM logits and computes loss on response tokens only.
122/// GPU-resident training state for NF4 QLoRA backward pass.
123///
124/// Holds per-layer activation snapshots and scratch buffers needed for
125/// activation checkpointing during NF4 backward.
126#[cfg(feature = "cuda")]
127pub(super) struct InstructGpuTrainingState {
128    /// Saved input to each block during forward [num_layers][max_seq_len * hidden_size]
129    layer_inputs: Vec<GpuBuffer<f32>>,
130    /// Final RMSNorm weight uploaded to GPU [hidden_size]
131    final_norm_weight: GpuBuffer<f32>,
132    /// Blocks output saved on GPU for final norm backward [max_seq_len * hidden_size]
133    blocks_output: GpuBuffer<f32>,
134    /// Gradient scratch buffer A [max_seq_len * hidden_size]
135    grad_buf_a: GpuBuffer<f32>,
136    /// Gradient scratch buffer B [max_seq_len * hidden_size]
137    grad_buf_b: GpuBuffer<f32>,
138    /// Gradient for final RMSNorm weight [hidden_size]
139    grad_final_norm_weight: GpuBuffer<f32>,
140    embed_transposed: GpuBuffer<f32>, // [hidden*vocab] lm_head forward
141    embed_original: GpuBuffer<f32>,   // [vocab*hidden] lm_head backward (KAIZEN-068)
142    /// GPU scratch for logits [max_seq_len * vocab_size]
143    logits_buf: GpuBuffer<f32>,
144    /// GPU scratch for grad_hidden [max_seq_len * hidden_size]
145    grad_hidden_buf: GpuBuffer<f32>,
146    /// KAIZEN-045: Pre-allocated scratch buffer for activation checkpointing in backward
147    output_scratch: GpuBuffer<f32>,
148    /// KAIZEN-045: Pre-allocated upload buffer for gradient H2D transfer in backward
149    grad_upload_buf: GpuBuffer<f32>,
150    /// KAIZEN-062: Pre-allocated forward ping-pong buffer A
151    fwd_scratch_a: GpuBuffer<f32>,
152    /// KAIZEN-062: Pre-allocated forward ping-pong buffer B
153    fwd_scratch_b: GpuBuffer<f32>,
154    /// KAIZEN-062: Pre-allocated lm_head hidden input buffer
155    lm_head_hidden_buf: GpuBuffer<f32>,
156    /// PMAT-464: Cached CUDA graph for forward pass replay.
157    forward_graph_exec: Option<trueno_gpu::driver::CudaGraphExec>,
158    graph_cached_seq_len: usize,
159    /// PMAT-488: Cached CUDA graph for backward pass replay.
160    backward_graph_state: Option<super::backward_graph::BackwardGraphState>,
161    /// PMAT-063: cuBLAS workspace buffer (must outlive CUDA graph)
162    cublas_workspace: Option<GpuBuffer<f32>>,
163    /// PMAT-483: Per-layer forward timing (microseconds per layer per step)
164    profiler_layer_fwd_us: Vec<u64>,
165    /// PMAT-483: Per-layer backward timing (microseconds per layer per step)
166    profiler_layer_bwd_us: Vec<u64>,
167    /// PMAT-483: Temporary layer start timestamp
168    profiler_layer_start: Option<std::time::Instant>,
169    /// PMAT-483/entrenar#328: Per-operation timing within layers (accumulated per step)
170    /// Index matches StepProfiler::OP_* constants. Reset each step.
171    profiler_op_us: [u64; 16],
172    /// Per-operation start timestamp
173    profiler_op_start: Option<std::time::Instant>,
174}
175
176pub struct InstructPipeline {
177    /// Base transformer model
178    pub model: Transformer,
179    /// LoRA adapters applied to Q/V attention projections
180    pub lora_layers: Vec<LoRALayer>,
181    /// Pipeline configuration
182    pub config: InstructConfig,
183    /// AdamW optimizer for trainable parameters
184    optimizer: AdamW,
185    /// Optional BPE tokenizer
186    tokenizer: Option<HfTokenizer>,
187    /// Path to base model (for checkpoint provenance)
188    model_dir: Option<PathBuf>,
189    /// PMAT-483: Per-step profiler for scientific training measurement.
190    /// Zero-overhead when disabled. Enable via --profile-interval N.
191    pub profiler: StepProfiler,
192    /// CUDA trainer for GPU memory management
193    #[cfg(feature = "cuda")]
194    cuda_trainer: Option<CudaTrainer>,
195    /// CUDA-accelerated transformer blocks -- one per layer
196    #[cfg(feature = "cuda")]
197    cuda_blocks: Option<Vec<CudaBlock>>,
198    /// Shared scratch buffers for NF4 forward pass
199    #[cfg(feature = "cuda")]
200    shared_scratch: Option<CudaBlockScratch>,
201    /// Count of GPU forward passes that produced NaN/Inf
202    #[cfg(feature = "cuda")]
203    #[allow(dead_code)]
204    cuda_nan_count: usize,
205    /// GPU training state for NF4 QLoRA backward pass
206    #[cfg(feature = "cuda")]
207    gpu_training: Option<InstructGpuTrainingState>,
208    /// Shared LoRA gradient workspace for NF4 QLoRA backward
209    #[cfg(feature = "cuda")]
210    cuda_lora_grad_workspace: Option<CudaLoraGradWorkspace>,
211    /// PMAT-477: Fused clip state -- zero D2H sync gradient clipping
212    #[cfg(feature = "cuda")]
213    lora_fused_clip: Option<crate::autograd::cuda_optim::FusedClipState>,
214    /// Per-layer LoRA optimizer states for NF4 QLoRA training
215    #[cfg(feature = "cuda")]
216    cuda_lora_optimizer_states: Option<Vec<GpuLoraOptimizerState>>,
217    /// NF4 LoRA optimizer step counter
218    #[cfg(feature = "cuda")]
219    nf4_lora_step: u32,
220    /// VRAM reservation guard (GPU-SHARE-002). Releases ledger entry on Drop.
221    #[cfg(feature = "cuda")]
222    #[allow(dead_code)]
223    vram_guard: Option<VramGuard>,
224    /// wgpu training pipeline (zero unsafe alternative to CUDA)
225    #[cfg(feature = "gpu")]
226    wgpu_training: Option<WgpuTrainingState>,
227}
228
229/// State for wgpu-based training pipeline (WgpuTrainingPipeline)
230#[cfg(feature = "gpu")]
231struct WgpuTrainingState {
232    /// GPU forward pass with persistent weight buffers + tiled GEMM
233    fwd: trueno::backends::gpu::WgslForwardPass,
234    cross_entropy: crate::autograd::wgpu_cross_entropy::WgslCrossEntropy,
235    trainer: crate::autograd::wgpu_training::WgpuTrainer,
236    // GPU buffers for logits, labels, losses, logsumexp
237    logits_buf: trueno::backends::gpu::wgpu::Buffer,
238    labels_buf: trueno::backends::gpu::wgpu::Buffer,
239    losses_buf: trueno::backends::gpu::wgpu::Buffer,
240    logsumexp_buf: trueno::backends::gpu::wgpu::Buffer,
241    // Precomputed lm_head GPU buffers
242    lm_head_gpu: trueno::backends::gpu::wgpu::Buffer,
243    lm_head_t_gpu: trueno::backends::gpu::wgpu::Buffer,
244    // Model config needed for forward pass
245    num_layers: usize,
246    hidden_dim: usize,
247    vocab_size: usize,
248}
249
250/// Configuration for autoregressive text generation.
251#[derive(Debug, Clone)]
252pub struct GenerateConfig {
253    /// Maximum number of new tokens to generate (default: 256)
254    pub max_new_tokens: usize,
255    /// Sampling temperature (0.0 = greedy/argmax, >0 = stochastic)
256    pub temperature: f32,
257    /// Top-k filtering (0 = disabled, >0 = keep only top-k logits)
258    pub top_k: usize,
259    /// Additional stop token IDs (generation stops on EOS or any of these)
260    pub stop_tokens: Vec<u32>,
261}
262
263/// Sample a token from logits with temperature and top-k filtering.
264fn sample_token(logits: &[f32], temperature: f32, top_k: usize) -> u32 {
265    if temperature <= 0.0 || top_k == 1 {
266        // Greedy: argmax
267        return logits
268            .iter()
269            .enumerate()
270            .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
271            .map_or(0, |(idx, _)| idx as u32);
272    }
273
274    // Temperature scaling
275    let scaled: Vec<f32> = logits.iter().map(|&l| l / temperature).collect();
276
277    // Top-k filtering
278    let mut indices_and_logits: Vec<(usize, f32)> = scaled.iter().copied().enumerate().collect();
279    indices_and_logits
280        .sort_unstable_by(|(_, a), (_, b)| b.partial_cmp(a).unwrap_or(std::cmp::Ordering::Equal));
281
282    let k = if top_k > 0 && top_k < indices_and_logits.len() {
283        top_k
284    } else {
285        indices_and_logits.len()
286    };
287    let top = &indices_and_logits[..k];
288
289    // Softmax over top-k
290    let max_logit = top[0].1;
291    let exps: Vec<f32> = top.iter().map(|(_, l)| (l - max_logit).exp()).collect();
292    let sum: f32 = exps.iter().sum();
293    let probs: Vec<f32> = exps.iter().map(|e| e / sum).collect();
294
295    // Sample from distribution (simple linear scan)
296    let r: f32 = simple_random();
297    let mut cumulative = 0.0;
298    for (i, &p) in probs.iter().enumerate() {
299        cumulative += p;
300        if r < cumulative {
301            return top[i].0 as u32;
302        }
303    }
304
305    // Fallback to top-1
306    top[0].0 as u32
307}
308
309/// Simple pseudo-random float in [0, 1) using thread-local state.
310/// Not cryptographically secure but sufficient for sampling.
311fn simple_random() -> f32 {
312    use std::cell::Cell;
313    thread_local! {
314        static STATE: Cell<u64> = Cell::new(
315            std::time::SystemTime::now()
316                .duration_since(std::time::UNIX_EPOCH)
317                .map(|d| d.as_nanos() as u64)
318                .unwrap_or(42)
319        );
320    }
321    STATE.with(|s| {
322        // xorshift64
323        let mut x = s.get();
324        x ^= x << 13;
325        x ^= x >> 7;
326        x ^= x << 17;
327        s.set(x);
328        (x >> 40) as f32 / (1u64 << 24) as f32
329    })
330}
entrenar/finetune/instruct_pipeline/mod.rs

entrenar/finetune/instruct_pipeline/
mod.rs
