wip

Author: tastelikefeet

cookbook/transformers/fsdp2.py

@@ -10,7 +10,7 @@
 from twinkle.preprocessor import SelfCognitionProcessor
 
 # Construct a device_mesh, dp=2
-device_mesh = DeviceMesh.from_sizes(dp_size=2)
+device_mesh = DeviceMesh.from_sizes(dp_size=8)
 # use torchrun mode
 twinkle.initialize(mode='local', global_device_mesh=device_mesh)
 

src/twinkle/model/megatron/megatron.py

@@ -245,7 +245,7 @@ def __init__(
 
     def _construct_default_optimizer_group(self):
         return MegatronOptimizerGroup(
-            loss_instance=CrossEntropyLoss(),
+            loss_instance=CrossEntropyLoss(reduction='sum'),
             template=Template(self.tokenizer_id),
             processor=InputProcessor(self.device_mesh, framework='megatron'),
             _device_mesh=self.device_mesh,