Merge pull request #2702 from AI-Hypercomputer:xfgu-multislice-rl

PiperOrigin-RevId: 833880778

Author: Google-ML-Automation

src/MaxText/configs/rl.yml

@@ -21,6 +21,8 @@ base_config: "base.yml"
 trainer_devices_fraction: 0.5
 sampler_devices_fraction: 0.5
 chips_per_vm: 4  # depends on hardware, for v5p this is 4
+num_trainer_slices: -1
+num_samplers_slices: -1
 
 # ====== Reproducibility ======
 data_shuffle_seed: 42

src/MaxText/configs/types.py

@@ -1223,6 +1223,8 @@ class RLHardware(BaseModel):
   sampler_devices_fraction: float = Field(0.5, description="Fraction of devices to use for the sampler.")
   chips_per_vm: int = Field(4, description="Number of accelerator chips per VM.")
   use_pathways: bool = Field(True, description="Whether to use Pathways for multihost orchestration.")
+  num_trainer_slices: int = Field(-1, description="Number of slices for the trainer.")
+  num_samplers_slices: int = Field(-1, description="Number of slices for the samplers.")
 
 
 class VLLM(BaseModel):