Merge pull request #3452 from AI-Hypercomputer:optimizer_ga_fix_2

Google-ML-Automation · Google-ML-Automation · commit 1a45bf0cd386 · 2026-03-20T11:15:24.000-07:00
PiperOrigin-RevId: 886886106
diff --git a/src/maxtext/trainers/post_train/distillation/distillation_utils.py b/src/maxtext/trainers/post_train/distillation/distillation_utils.py
@@ -301,6 +301,9 @@ def compute_loss(
         "distill/teacher_loss": teacher_hard_loss,
         "distill/out_proj_feature_loss": feature_loss,
         "distill/total_loss": total_loss,
+        "distill/temperature": self.temperature,
+        "distill/alpha": self.alpha,
+        "distill/beta_feature": self.beta_feature,
     }
     return total_loss, metrics
 
diff --git a/src/maxtext/trainers/post_train/distillation/train_distill.py b/src/maxtext/trainers/post_train/distillation/train_distill.py
@@ -480,8 +480,7 @@ def labels_fn(targets, targets_segmentation=None, **kwargs):
   )
 
   # 4. Optimizer & Config
-  total_updates = student_config.steps // student_config.gradient_accumulation_steps
-  optimizer = get_distillation_optimizer(student_config, total_updates)
+  optimizer = get_distillation_optimizer(student_config, student_config.steps)
 
   checkpointing_options = checkpoint.CheckpointManagerOptions(
       save_interval_steps=student_config.checkpoint_period,
diff --git a/tests/post_training/unit/train_distill_test.py b/tests/post_training/unit/train_distill_test.py
@@ -399,6 +399,9 @@ def _test_monitored_strategy(self, sft_mode: bool):
         "distill/teacher_loss",
         "distill/out_proj_feature_loss",
         "distill/total_loss",
+        "distill/temperature",
+        "distill/alpha",
+        "distill/beta_feature",
     ]
     for key in expected_keys:
       self.assertIn(key, metrics)

Original file line number	Diff line number	Diff line change
`@@ -301,6 +301,9 @@ def compute_loss(`
`301`	`301`	`"distill/teacher_loss": teacher_hard_loss,`
`302`	`302`	`"distill/out_proj_feature_loss": feature_loss,`
`303`	`303`	`"distill/total_loss": total_loss,`
	`304`	`+ "distill/temperature": self.temperature,`
	`305`	`+ "distill/alpha": self.alpha,`
	`306`	`+ "distill/beta_feature": self.beta_feature,`
`304`	`307`	`}`
`305`	`308`	`return total_loss, metrics`
`306`	`309`
Original file line number	Diff line number	Diff line change
`@@ -480,8 +480,7 @@ def labels_fn(targets, targets_segmentation=None, **kwargs):`
`480`	`480`	`)`
`481`	`481`
`482`	`482`	`# 4. Optimizer & Config`
`483`		`- total_updates = student_config.steps // student_config.gradient_accumulation_steps`
`484`		`- optimizer = get_distillation_optimizer(student_config, total_updates)`
	`483`	`+ optimizer = get_distillation_optimizer(student_config, student_config.steps)`
`485`	`484`
`486`	`485`	`checkpointing_options = checkpoint.CheckpointManagerOptions(`
`487`	`486`	`save_interval_steps=student_config.checkpoint_period,`