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

feat(ml): TFT batched gradient norm — N GPU syncs → 1

Browse Source 
Replace per-param gradient norm extraction loop with batched
Tensor::stack pattern. Single GPU→CPU sync instead of one per
parameter tensor.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-02 17:30:10 +01:00
parent 7d2b56a3ad
commit 4cf181c8fb
1 changed files with 22 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
crates/ml/src/tft/training.rs
View File

@@ -479,18 +479,28 @@ impl TFTTrainer {
let varmap_data = self.model.varmap().data().lock().map_err(|e| {
MLError::TrainingError(format!("Failed to lock VarMap: {}", e))
})?;
let mut total_norm_sq = 0.0_f64;
for (_name, var) in varmap_data.iter() {
if let Some(grad) = grads.get(var.as_tensor()) {
let norm_sq = grad
.sqr()
.and_then(|t| t.sum_all())
.and_then(|t| t.to_dtype(DType::F64))
.and_then(|t| t.to_scalar::<f64>())
.unwrap_or(0.0);
total_norm_sq += norm_sq;
}
}
let norm_parts: Vec<Tensor> = varmap_data
.iter()
.filter_map(|(_name, var)| {
grads
.get(var.as_tensor())
.and_then(|g| g.sqr().and_then(|s| s.sum_all()).ok())
})
.collect();
let total_norm_sq: f64 = if !norm_parts.is_empty() {
let stacked = Tensor::stack(&norm_parts, 0).map_err(|e| {
MLError::TrainingError(format!("TFT grad norm stack: {e}"))
})?;
stacked
.sum_all()
.and_then(|s| s.to_scalar::<f64>())
.map_err(|e| {
MLError::TrainingError(format!("TFT grad norm: {e}"))
})?
} else {
0.0
};
drop(varmap_data);
let grad_norm = total_norm_sq.sqrt();
self.last_gradient_norm = grad_norm;