gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(ml): PPO mixed precision auto-detection — BF16 on A100/H100

Browse Source 
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-01 20:43:27 +01:00
parent 5d390037e5
commit 6deaf86643
1 changed files with 15 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
crates/ml/src/trainers/ppo.rs
View File

@@ -305,6 +305,21 @@ impl PpoTrainer {
config.gae_config.gamma = hyperparams.gamma as f32;
config.gae_config.lambda = hyperparams.gae_lambda;
// Auto-detect mixed precision capability based on GPU architecture
if device.is_cuda() {
match crate::memory_optimization::auto_batch_size::detect_gpu_memory() {
Ok((_total, _free, ref name)) => {
let detected = crate::dqn::mixed_precision::detect_from_gpu_name(name);
match &detected {
Some(c) => info!("PPO mixed precision: {:?} enabled (GPU: {})", c.dtype, name),
None => info!("PPO mixed precision: disabled (GPU: {})", name),
}
config.mixed_precision = detected;
}
Err(_) => {}
}
}
// Create PPO model with specified device (GPU or CPU)
let model = PPO::with_device(config, device.clone())?;