fix

Author: tastelikefeet

cookbook/legacy/grpo/gsm8k.py

@@ -237,7 +237,7 @@ def main():
             name='sampler',
             ranks=list(range(MODEL_GPUS, NUM_GPUS)),
             device_type='GPU',
-            gpus_per_worker=1,
+            gpus_per_worker=4,
         ),
     ]
     if USE_MEGATRON:
@@ -249,7 +249,7 @@ def main():
             world_size=MODEL_GPUS, dp_size=MODEL_GPUS,
         )
     sampler_mesh = DeviceMesh.from_sizes(
-        world_size=SAMPLER_GPUS, dp_size=4
+        world_size=SAMPLER_GPUS, tp_size=4
     )
     twinkle.initialize(
         mode='ray',

src/twinkle/checkpoint_engine/manager.py

@@ -81,7 +81,7 @@ def sync_weights(self, merge_and_sync=True):
         model_metadata = self.model.prepare_checkpoint_engine([True] + [False]*(self.model.device_mesh.world_size -1))
         self.sampler.prepare_checkpoint_engine(False)
         model_kwargs, sampler_kwargs = self.backend_cls.build_topology(
-            self.model.device_mesh.world_size, self.sampler.device_mesh.world_size, [model_metadata],
+            self.model.device_mesh.world_size, self.sampler.device_mesh.data_world_size, [model_metadata],
         )
         # Launch both init calls concurrently — TCPStore server (model rank 0)
         # blocks until all clients (sampler ranks) connect, so these MUST NOT