added optional keys to config to make the code cleaner

Author: ajkv-google

src/maxtext/configs/types.py

@@ -1052,6 +1052,13 @@ class Distillation(BaseModel):
   # --- Loss Params ---
   distill_alpha: float = Field(0.5, description="Weight for the distillation loss component.")
   distill_temperature: float = Field(1.0, description="Temperature for distillation softening.")
+  
+  # --- Teacher topk distillation ---
+  teacher_logits_optional_keys: list[str] = Field(
+      default=["inputs_position", "inputs_segmentation", "targets_segmentation", "targets"],
+      description="Optional keys to save from teacher logits"
+  )
+  
 
 
 class TrainingLoop(BaseModel):
@@ -1809,7 +1816,6 @@ class MaxTextConfig(
     # Reinforcement Learning
     RLHardware,
     VLLM,
-    RL,
     RLDataset,
     RLEvaluation,
     Reward,

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

@@ -72,7 +72,7 @@ def generate_and_save_data(config, k_val):
     gathered_idx = multihost_utils.process_allgather(top_k_idx, tiled=True)
     gathered_tokens = multihost_utils.process_allgather(tokens, tiled=True)
 
-    optional_keys = ["inputs_position", "inputs_segmentation", "targets_segmentation", "targets"]
+    optional_keys = config.teacher_logits_optional_keys
     gathered_optionals = {
         key: multihost_utils.process_allgather(batch[key], tiled=True) for key in optional_keys if key in batch
     }