fix

Author: tastelikefeet

cookbook/megatron/ddp.py

@@ -11,7 +11,7 @@
 from twinkle.model import MegatronModel
 from twinkle.preprocessor import SelfCognitionProcessor
 
-if Platform.get_rank() == 0:
+if Platform.get_rank() == 0 and os.environ.get('SWANLAB_API_KEY'):
     # rank0 recording
     import swanlab
     swanlab.login(api_key=os.environ['SWANLAB_API_KEY'], save=True)
@@ -84,7 +84,7 @@ def train():
         if step % 5 == 0:
             # Print metric
             metric = model.calculate_metric(is_training=True)
-            if Platform.get_rank() == 0:
+            if Platform.get_rank() == 0 and os.environ.get('SWANLAB_API_KEY'):
                 swanlab.log(metric)
             logger.info(f'Current is step {step} of {len(dataloader)}, metric: {metric}')
         if step > 0 and step % 20 == 0:

cookbook/megatron/ddp.sh

@@ -0,0 +1 @@
+CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python3 ddp.py

cookbook/megatron/ddp_moe.py

@@ -11,7 +11,7 @@
 from twinkle.model import MegatronModel
 from twinkle.preprocessor import SelfCognitionProcessor
 
-if Platform.get_rank() == 0:
+if Platform.get_rank() == 0 and os.environ.get('SWANLAB_API_KEY'):
     # rank0 recording
     import swanlab
     swanlab.login(api_key=os.environ['SWANLAB_API_KEY'], save=True)
@@ -22,7 +22,7 @@
 
 
 # Construct a device_mesh, tp=pp=cp=ep=2, dp=1
-device_mesh = DeviceMesh.from_sizes(dp_size=1, tp_size=2, pp_size=2, cp_size=2, ep_size=2, sequence_parallel=True)
+device_mesh = DeviceMesh.from_sizes(dp_size=1, tp_size=2, pp_size=2, cp_size=2, ep_size=2)
 # use torchrun mode
 twinkle.initialize(mode='local', global_device_mesh=device_mesh)
 
@@ -74,8 +74,7 @@ def train():
     logger.info(model.get_train_configs())
     logger.info(f'Total steps: {len(dataloader)}')
     loss_metric = 99.0
-    # lora: 10G * 8
-    # full: 40G * 8
+    # lora: 23G * 8
     for step, batch in enumerate(dataloader):
         # Do forward and backward
         model.forward_backward(inputs=batch)
@@ -84,7 +83,7 @@ def train():
         if step % 5 == 0:
             # Print metric
             metric = model.calculate_metric(is_training=True)
-            if Platform.get_rank() == 0:
+            if Platform.get_rank() == 0 and os.environ.get('SWANLAB_API_KEY'):
                 swanlab.log(metric)
             logger.info(f'Current is step {step} of {len(dataloader)}, metric: {metric}')
         if step > 0 and step % 20 == 0:

cookbook/megatron/ddp_moe.sh

@@ -0,0 +1 @@
+CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python3 ddp_moe.py

cookbook/transformers/fsdp2.py

@@ -11,7 +11,7 @@
 from twinkle.model import TransformersModel
 from twinkle.preprocessor import SelfCognitionProcessor
 
-if Platform.get_rank() == 0:
+if Platform.get_rank() == 0 and os.environ.get('SWANLAB_API_KEY'):
     # rank0 recording
     import swanlab
     swanlab.login(api_key=os.environ['SWANLAB_API_KEY'], save=True)
@@ -21,8 +21,8 @@
     )
 
 
-# Construct a device_mesh, fsdp=2, dp=2
-device_mesh = DeviceMesh.from_sizes(fsdp_size=4)
+# Construct a device_mesh, fsdp=4, dp=2
+device_mesh = DeviceMesh.from_sizes(fsdp_size=4, dp_size=2)
 # use torchrun mode
 twinkle.initialize(mode='local', global_device_mesh=device_mesh)
 
@@ -53,7 +53,7 @@ def train():
     # Encode dataset
     dataset.encode()
     # Global batch size = 4, for GPUs, so 1 sample per GPU
-    dataloader = DataLoader(dataset=dataset, batch_size=4)
+    dataloader = DataLoader(dataset=dataset, batch_size=8)
     # Use a TransformersModel
     model = TransformersModel(model_id='ms://Qwen/Qwen2.5-7B-Instruct')
 
@@ -64,7 +64,8 @@ def train():
     )
 
     # Add a lora to model, with name `default`
-    model.add_adapter_to_model('default', lora_config, gradient_accumulation_steps=4)
+    # Comment this to use full-parameter training
+    model.add_adapter_to_model('default', lora_config, gradient_accumulation_steps=2)
     # Add Optimizer for lora `default`
     model.set_optimizer(optimizer_cls='AdamW', lr=1e-4)
     # Add LRScheduler for lora `default`
@@ -84,7 +85,7 @@ def train():
         if step % 20 == 0:
             # Print metric
             metric = model.calculate_metric(is_training=True)
-            if Platform.get_rank() == 0:
+            if Platform.get_rank() == 0 and os.environ.get('SWANLAB_API_KEY'):
                 swanlab.log(metric)
             logger.info(f'Current is step {step} of {len(dataloader)}, metric: {metric}')
         if step > 0 and step % 40 == 0:

cookbook/transformers/fsdp2.sh

@@ -0,0 +1 @@
+CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python3 fsdp2.py

cookbook/transformers/fsdp2_moe.py

@@ -11,7 +11,7 @@
 from twinkle.model import TransformersModel
 from twinkle.preprocessor import SelfCognitionProcessor
 
-if Platform.get_rank() == 0:
+if Platform.get_rank() == 0 and os.environ.get('SWANLAB_API_KEY'):
     # rank0 recording
     import swanlab
     swanlab.login(api_key=os.environ['SWANLAB_API_KEY'], save=True)
@@ -65,6 +65,7 @@ def train():
     )
 
     # Add a lora to model, with name `default`
+    # Comment this to use full-parameter training
     model.add_adapter_to_model('default', lora_config, gradient_accumulation_steps=2)
     # Add Optimizer for lora `default`
     model.set_optimizer(optimizer_cls='AdamW', lr=1e-4)
@@ -75,8 +76,7 @@ def train():
     logger.info(model.get_train_configs())
     logger.info(f'Total steps: {len(dataloader)}')
     loss_metric = 99.0
-    # lora: 18G * 4
-    # full: 50G * 4
+    # lora: 34G * 8
     for step, batch in enumerate(dataloader):
         # Do forward and backward
         model.forward_backward(inputs=batch)
@@ -85,7 +85,7 @@ def train():
         if step % 20 == 0:
             # Print metric
             metric = model.calculate_metric(is_training=True)
-            if Platform.get_rank() == 0:
+            if Platform.get_rank() == 0 and os.environ.get('SWANLAB_API_KEY'):
                 swanlab.log(metric)
             logger.info(f'Current is step {step} of {len(dataloader)}, metric: {metric}')
         if step > 0 and step % 40 == 0:

cookbook/transformers/fsdp2_moe.sh

@@ -0,0 +1 @@
+CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python3 fsdp2_moe.py