Removed redundant offline_distillation flag and relied on offline_data_dir to know when to run offfline vs online distillation

Author: ajkv-google

src/maxtext/configs/types.py

@@ -1082,10 +1082,7 @@ class Distillation(BaseModel):
       description="Overrides specific to the Teacher model (e.g., {'num_query_heads': 64}).",
   )
 
-  # --- Offline Distillation Fields ---
-  offline_distillation: bool = Field(
-      False, description="If True, enables offline distillation using pre-generated teacher data."
-  )
+  # --- Offline Distillation Field ---
   offline_data_dir: Optional[str] = Field(
       None, description="GCS or local path to the pre-generated ArrayRecord teacher data."
   )

src/maxtext/trainers/post_train/distillation/train_distill.py

@@ -629,12 +629,14 @@ def main(argv: Sequence[str]) -> None:
   student_overrides = global_config.student_overrides
   student_config = pyconfig.initialize(argv, **student_overrides)
 
+  is_offline = bool(global_config.offline_data_dir)
+
   # 3. Initialize TEACHER Config
   # We isolate the Teacher from Student CLI arguments (like pruning params).
   teacher_overrides = global_config.teacher_overrides
 
   # Ensure load_parameters_path is set in overrides
-  if not global_config.offline_distillation and not teacher_overrides.get("load_parameters_path"):
+  if not is_offline and not teacher_overrides.get("load_parameters_path"):
     raise ValueError(
         "Teacher model path is missing! You must provide 'teacher_overrides.load_parameters_path' "
         "in your config or arguments."
@@ -646,7 +648,7 @@ def main(argv: Sequence[str]) -> None:
   teacher_config = pyconfig.initialize(teacher_argv, **teacher_overrides)
 
   # 4. Run Training
-  train_distill(student_config, teacher_config, global_config.offline_distillation, global_config.offline_data_dir)
+  train_distill(student_config, teacher_config, is_offline, global_config.offline_data_dir)
 
 
 if __name__ == "__main__":