Introduce Multislice RL

Author: xuefgu

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

@@ -1218,6 +1218,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):