apr_cli/
train_commands.rs

1
2/// Training pipeline subcommands (forjar-style plan/apply).
3///
4/// Thin CLI wrappers around entrenar's training plan/apply infrastructure.
5#[derive(Subcommand, Debug)]
6pub enum TrainCommands {
7    /// Generate a training plan without touching the GPU.
8    ///
9    /// Validates data quality, checks model compatibility, builds HPO search space,
10    /// estimates resource usage, and runs pre-flight checks. Outputs a serializable
11    /// plan manifest (text, JSON, or YAML).
12    ///
13    /// Analogous to `forjar plan` — shows what will happen before committing GPU time.
14    Plan {
15        /// Path to training data (JSONL). Only read by `apr finetune --task classify`.
16        #[arg(long, value_name = "FILE")]
17        data: Option<PathBuf>,
18        /// Model size: "0.5B", "9B", "7B", "13B"
19        #[arg(long, default_value = "0.5B")]
20        model_size: String,
21        /// Path to model weights directory
22        #[arg(long, value_name = "DIR")]
23        model_path: Option<PathBuf>,
24        /// Number of output classes
25        #[arg(long, default_value = "5")]
26        num_classes: usize,
27        /// Task type: pretrain (causal LM). Classification fine-tuning is
28        /// `apr finetune --task classify`, not this command.
29        #[arg(long, default_value = "pretrain")]
30        task: String,
31        /// YAML training config (for --task pretrain)
32        #[arg(long, value_name = "FILE")]
33        config: Option<PathBuf>,
34        /// Output directory for checkpoints
35        #[arg(short, long, default_value = "/tmp/training-output")]
36        output: PathBuf,
37        /// HPO strategy: tpe, grid, random, manual
38        #[arg(long, default_value = "tpe")]
39        strategy: String,
40        /// HPO budget (number of trials)
41        #[arg(long, default_value = "20")]
42        budget: usize,
43        /// Scout mode: 1 epoch per trial for fast exploration
44        #[arg(long)]
45        scout: bool,
46        /// Maximum epochs per trial
47        #[arg(long, default_value = "3")]
48        max_epochs: usize,
49        /// Manual learning rate (only used with --strategy manual)
50        #[arg(long)]
51        learning_rate: Option<f32>,
52        /// Manual LoRA rank (only used with --strategy manual)
53        #[arg(long)]
54        lora_rank: Option<usize>,
55        /// Manual batch size (only used with --strategy manual)
56        #[arg(long)]
57        batch_size: Option<usize>,
58        /// Validation data file (JSONL)
59        #[arg(long, value_name = "FILE")]
60        val_data: Option<PathBuf>,
61        /// Test data file (JSONL)
62        #[arg(long, value_name = "FILE")]
63        test_data: Option<PathBuf>,
64        /// Output format: text, json, yaml
65        #[arg(long, default_value = "text")]
66        format: String,
67    },
68
69    /// Execute a training plan (allocate GPU, run trials).
70    ///
71    /// Reads a previously generated plan (YAML/JSON) and executes it:
72    /// - Manual strategy: single training run with specified hyperparameters
73    /// - HPO strategy: multiple trials with automatic hyperparameter tuning
74    ///
75    /// Analogous to `forjar apply` — commits resources and executes the plan.
76    Apply {
77        /// Path to a saved plan file (YAML or JSON from `apr train plan`)
78        #[arg(long, value_name = "FILE")]
79        plan: Option<PathBuf>,
80
81        /// YAML training config (for --task pretrain)
82        #[arg(long, value_name = "FILE")]
83        config: Option<PathBuf>,
84
85        /// Task type: pretrain (causal LM). Classification fine-tuning is
86        /// `apr finetune --task classify`, not this command.
87        #[arg(long, default_value = "pretrain")]
88        task: String,
89
90        // ── Inline plan params (used when no --plan file is given) ─────
91        /// Path to training data (JSONL)
92        #[arg(long, value_name = "FILE")]
93        data: Option<PathBuf>,
94        /// Model size: "0.5B", "9B", "7B", "13B"
95        #[arg(long, default_value = "0.5B")]
96        model_size: String,
97        /// Path to model weights directory
98        #[arg(long, value_name = "DIR")]
99        model_path: Option<PathBuf>,
100        /// Number of output classes
101        #[arg(long, default_value = "5")]
102        num_classes: usize,
103        /// Output directory for checkpoints and leaderboard.
104        ///
105        /// When given it OVERRIDES `training.output_dir` in the YAML config.
106        /// When omitted, the config's `training.output_dir` is used (default
107        /// `./checkpoints`). The directory is created if it does not exist.
108        #[arg(short, long, value_name = "DIR")]
109        output: Option<PathBuf>,
110        /// HPO strategy: tpe, grid, random, manual
111        #[arg(long, default_value = "tpe")]
112        strategy: String,
113        /// HPO budget (number of trials)
114        #[arg(long, default_value = "20")]
115        budget: usize,
116        /// Scout mode: 1 epoch per trial
117        #[arg(long)]
118        scout: bool,
119        /// Maximum epochs per trial
120        #[arg(long, default_value = "3")]
121        max_epochs: usize,
122        /// Manual learning rate (only used with --strategy manual)
123        #[arg(long)]
124        learning_rate: Option<f32>,
125        /// Manual LoRA rank (only used with --strategy manual)
126        #[arg(long)]
127        lora_rank: Option<usize>,
128        /// Manual batch size (only used with --strategy manual)
129        #[arg(long)]
130        batch_size: Option<usize>,
131
132        // ── Distributed training params (tickets #131-#140, aprender #393) ──
133        /// Enable distributed data-parallel training
134        #[arg(long)]
135        distributed: bool,
136        /// Total number of workers (default: auto-detect GPUs)
137        #[arg(long, value_name = "N")]
138        world_size: Option<usize>,
139        /// This worker's global rank (default: 0 = coordinator)
140        #[arg(long, value_name = "N")]
141        rank: Option<usize>,
142        /// Coordinator address for distributed training (default: 0.0.0.0:9000)
143        #[arg(long, value_name = "HOST:PORT")]
144        coordinator_addr: Option<String>,
145
146        // ── Reproducibility params (R-084 C-DETERM-001) ──
147        /// Enable bitwise deterministic training (CUBLAS_WORKSPACE_CONFIG, cuDNN deterministic)
148        #[arg(long)]
149        deterministic: bool,
150        /// Random seed for reproducibility (default: from YAML or 42)
151        #[arg(long, value_name = "N")]
152        seed: Option<u64>,
153
154        // ── Profiling params (PMAT-486) ──
155        /// Enable StepProfiler for per-phase wall-clock timing (KAIZEN-047)
156        #[arg(long)]
157        profile: bool,
158        /// StepProfiler report interval (every N steps, default: 50)
159        #[arg(long, value_name = "N", default_value = "50")]
160        profile_interval: usize,
161    },
162
163    /// Watch a training run with automatic restart on crash and hang detection.
164    ///
165    /// Monitors a running or to-be-started training process:
166    /// - Detects crashes (SIGABRT, SIGSEGV, OOM) and restarts with backoff
167    /// - Detects hangs via heartbeat/training_state.json staleness
168    /// - Captures GPU state and crash diagnostics
169    /// - Auto-enables CUDA_LAUNCH_BLOCKING on async crash pattern
170    ///
171    /// Sovereign Rust replacement for train-guard.sh.
172    Watch {
173        /// YAML training config to run and watch
174        #[arg(long, value_name = "FILE")]
175        config: PathBuf,
176
177        /// Maximum number of restart attempts
178        #[arg(long, default_value = "5")]
179        max_restarts: usize,
180
181        /// Heartbeat staleness threshold in seconds
182        #[arg(long, default_value = "300")]
183        heartbeat_timeout: u64,
184
185        /// Initial backoff delay in seconds
186        #[arg(long, default_value = "30")]
187        backoff_initial: u64,
188
189        /// Maximum backoff delay in seconds
190        #[arg(long, default_value = "600")]
191        backoff_max: u64,
192    },
193
194    /// Generate hyperparameter sweep configs from a base YAML.
195    ///
196    /// Creates N training configs with varied hyperparameters using grid
197    /// or random search. Each config is a complete YAML that can be
198    /// passed to `apr train apply --task pretrain --config <file>`.
199    ///
200    /// Sovereign Rust replacement for hyperparam-sweep.py.
201    Sweep {
202        /// Base YAML training config to sweep from
203        #[arg(long, value_name = "FILE")]
204        config: PathBuf,
205
206        /// Search strategy: grid or random
207        #[arg(long, default_value = "random")]
208        strategy: String,
209
210        /// Number of configs to generate (random) or max combinations (grid)
211        #[arg(long, default_value = "10")]
212        num_configs: usize,
213
214        /// Output directory for generated configs
215        #[arg(long, default_value = "sweeps/")]
216        output_dir: PathBuf,
217
218        /// Seed for random search reproducibility
219        #[arg(long, default_value = "42")]
220        seed: u64,
221    },
222
223    /// Run successive halving HPO on sweep configs (C-HPO-001).
224    ///
225    /// Takes a directory of sweep configs (from `apr train sweep`), runs each
226    /// for `--steps-per-round` steps, kills the worst half by val_ppl, doubles
227    /// steps, and repeats for `--rounds` rounds. Reports the winner with
228    /// μTransfer-scaled LR for the target model width.
229    ///
230    /// References: Hyperband (Li et al. 2018, arXiv:1603.06560),
231    /// μTransfer (Yang et al. 2022, arXiv:2203.03466).
232    Halving {
233        /// Directory containing sweep-*.yaml configs (from `apr train sweep`)
234        #[arg(long, value_name = "DIR")]
235        sweep_dir: PathBuf,
236
237        /// Number of halving rounds (default: 3)
238        #[arg(long, default_value = "3")]
239        rounds: usize,
240
241        /// Training steps in first round (doubles each round)
242        #[arg(long, default_value = "500")]
243        steps_per_round: usize,
244
245        /// Proxy model hidden_size (for μTransfer scaling)
246        #[arg(long, default_value = "512")]
247        source_width: usize,
248
249        /// Target model hidden_size (for μTransfer scaling)
250        #[arg(long, default_value = "1024")]
251        target_width: usize,
252
253        /// Output JSON file for results
254        #[arg(long, default_value = "sweeps/hpo-results.json")]
255        output: PathBuf,
256    },
257
258    /// Archive a checkpoint into a release bundle.
259    ///
260    /// Packages model weights, config, training state, and metadata
261    /// into a self-contained directory with integrity manifest.
262    Archive {
263        /// Path to checkpoint directory
264        #[arg(value_name = "CHECKPOINT_DIR")]
265        checkpoint_dir: PathBuf,
266
267        /// Output archive directory
268        #[arg(short, long, value_name = "DIR")]
269        output: PathBuf,
270
271        /// Release version tag (e.g., "v1.0")
272        #[arg(long = "release-version")]
273        release_version: Option<String>,
274
275        /// Release notes
276        #[arg(long)]
277        notes: Option<String>,
278    },
279
280    /// Submit multi-adapter training jobs to a cluster (GPU-SHARE Phase 3).
281    ///
282    /// Reads a cluster.yaml config, places adapter jobs across nodes using
283    /// the greedy placement algorithm, and generates launch commands.
284    Submit {
285        /// Path to cluster config YAML
286        #[arg(long, value_name = "FILE")]
287        cluster: PathBuf,
288
289        /// Model checkpoint path (.apr)
290        #[arg(long, value_name = "FILE")]
291        model: PathBuf,
292
293        /// Adapter specs: DATA:CHECKPOINT pairs (one per adapter)
294        #[arg(long = "adapter", value_name = "DATA:CHECKPOINT")]
295        adapters: Vec<String>,
296
297        /// LoRA rank
298        #[arg(long, default_value = "16")]
299        rank: u32,
300
301        /// Number of training epochs
302        #[arg(long, default_value = "3")]
303        epochs: u32,
304
305        /// Estimated VRAM budget per adapter (MB)
306        #[arg(long, default_value = "6000")]
307        budget_mb: u64,
308
309        /// Dry run: show placement and commands without executing
310        #[arg(long)]
311        dry_run: bool,
312    },
313
314    /// Show cluster status: nodes, GPUs, adapter capacity (GPU-SHARE Phase 3).
315    ///
316    /// Reads a cluster.yaml config and displays node health, VRAM availability,
317    /// and adapter placement capacity.
318    ClusterStatus {
319        /// Path to cluster config YAML
320        #[arg(long, value_name = "FILE")]
321        cluster: PathBuf,
322    },
323}
apr_cli/train_commands.rs

apr_cli/
train_commands.rs
