wip

Author: tastelikefeet

cookbook/rl/dpo_lora.py

@@ -58,7 +58,7 @@
 from twinkle.dataset import Dataset, DatasetMeta
 from twinkle.loss import DPOLoss
 from twinkle.metric import DPOMetric
-from twinkle.model import MegatronModel
+from twinkle.model import MultiLoraMegatronModel
 from twinkle.preprocessor import EmojiDPOProcessor
 from twinkle.processor import InputProcessor
 
@@ -68,7 +68,7 @@
 MODEL_ID = os.environ.get('MODEL_ID', 'ms://Qwen/Qwen2.5-7B-Instruct')
 DATASET_ID = os.environ.get('DATASET_ID', 'ms://hjh0119/shareAI-Llama3-DPO-zh-en-emoji')
 
-MODEL_GPUS = int(os.environ.get('MODEL_GPUS', 2))
+MODEL_GPUS = int(os.environ.get('MODEL_GPUS', 8))
 
 BATCH_SIZE = int(os.environ.get('BATCH_SIZE', 2))  # Number of preference pairs
 MICRO_BATCH_SIZE = int(os.environ.get('MICRO_BATCH_SIZE', 2))
@@ -137,7 +137,7 @@ def main():
         DeviceGroup(name='policy', ranks=list(range(MODEL_GPUS)), device_type='GPU'),
     ]
 
-    policy_mesh = DeviceMesh.from_sizes(world_size=MODEL_GPUS, dp_size=MODEL_GPUS)
+    policy_mesh = DeviceMesh.from_sizes(world_size=MODEL_GPUS, dp_size=1, pp_size=2, cp_size=2, tp_size=2)
     twinkle.initialize(mode='ray', nproc_per_node=8, groups=device_groups)
 
     # ── DataLoader Setup ──────────────────────────────────────────────────────
@@ -157,15 +157,17 @@ def main():
         lora_dropout=0.05,
     )
 
-    policy_model = MegatronModel(
+    policy_model = MultiLoraMegatronModel(
         model_id=MODEL_ID,
         device_mesh=policy_mesh,
         remote_group='policy',
     )
     MAX_STEPS = len(dataloader)
     policy_model.add_adapter_to_model(ADAPTER_NAME, lora_config, gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS)
-    policy_model.set_optimizer('default', lr=LEARNING_RATE, weight_decay=0.01)
-    policy_model.set_lr_scheduler('default', lr_decay_steps=MAX_STEPS)
+    # policy_model.set_optimizer('AdamW', lr=LEARNING_RATE, weight_decay=0.01, adapter_name=ADAPTER_NAME)
+    # policy_model.set_lr_scheduler('CosineAnnealingLR', T_max=MAX_STEPS, adapter_name=ADAPTER_NAME)
+    policy_model.set_optimizer('default', lr=LEARNING_RATE, weight_decay=0.01, adapter_name=ADAPTER_NAME)
+    policy_model.set_lr_scheduler('default', lr_decay_steps=MAX_STEPS, adapter_name=ADAPTER_NAME)
 
     # Set up loss function and metrics
     loss_fn = DPOLoss(
@@ -174,10 +176,10 @@ def main():
         reference_free=False,  # We use base model as reference via disable_lora=True
         sft_weight=SFT_WEIGHT,
     )
-    policy_model.set_loss(loss_fn)
-    policy_model.add_metric(DPOMetric, beta=DPO_BETA)
-    policy_model.set_processor(InputProcessor)
-    policy_model.set_template('Template', model_id=MODEL_ID)
+    policy_model.set_loss(loss_fn, adapter_name=ADAPTER_NAME)
+    policy_model.add_metric(DPOMetric, beta=DPO_BETA, adapter_name=ADAPTER_NAME)
+    policy_model.set_processor(InputProcessor, adapter_name=ADAPTER_NAME)
+    policy_model.set_template('Template', model_id=MODEL_ID, adapter_name=ADAPTER_NAME)
 
     optim_step = 0
     logger.info(get_device_placement())
@@ -191,32 +193,32 @@ def main():
 
         # Get reference outputs using base model (without LoRA adapter)
         # disable_lora=True tells the model to skip LoRA and use base weights
-        ref_outputs = policy_model.forward_only(inputs=dpo_batch, micro_batch_size=2, disable_lora=True)
-
+        ref_outputs = policy_model.forward_only(inputs=dpo_batch, micro_batch_size=2, disable_lora=True, adapter_name=ADAPTER_NAME)
         # Forward-backward pass with DPO loss (using LoRA adapter)
         # ref_outputs is passed to loss which extracts logps internally
         policy_model.forward_backward(
             inputs=dpo_batch,
             ref_outputs=ref_outputs,
             micro_batch_size=2,
+            adapter_name=ADAPTER_NAME
         )
 
         # Gradient clipping and optimizer step
-        policy_model.clip_grad_and_step()
+        policy_model.clip_grad_and_step(adapter_name=ADAPTER_NAME)
         optim_step += 1
 
         # Logging
         if optim_step % 1 == 0:
-            metrics = policy_model.calculate_metric(is_training=True)
+            metrics = policy_model.calculate_metric(is_training=True, adapter_name=ADAPTER_NAME)
             logger.info(f'[Step {optim_step}/{MAX_STEPS}] {metrics}')
 
         # Checkpointing
         if optim_step % SAVE_STEPS == 0:
-            policy_model.save(f'dpo-lora-checkpoint-{optim_step}')
+            policy_model.save(f'dpo-lora-checkpoint-{optim_step}', adapter_name=ADAPTER_NAME)
 
     # ── Save Final Checkpoint ─────────────────────────────────────────────────
     logger.info(f'Training completed. Total steps: {optim_step}')
-    policy_model.save('dpo-lora-final-checkpoint')
+    policy_model.save('dpo-lora-final-checkpoint', adapter_name=ADAPTER_NAME)
 
 
 if __name__ == '__main__':

src/twinkle/infra/collectors.py

@@ -61,7 +61,7 @@ def _pad_and_stack_tensors(tensors: List['torch.Tensor'], pad_value: float = -20
         raise ValueError('Empty tensor list')
 
     if len(tensors) == 1:
-        return tensors[0].unsqueeze(0)
+        return tensors[0]
 
     max_ndim = max(t.ndim for t in tensors)
     expanded_tensors = []

src/twinkle/model/megatron/megatron.py

@@ -400,7 +400,9 @@ def forward_backward(self,
                 seq_length = original_seq_length + (divisor - original_seq_length % divisor)
             else:
                 seq_length = original_seq_length
-
+        
+        if 'ref_outputs' in kwargs:
+            breakpoint()
         num_microbatches = len(inputs)
         loss_extra_kwargs_per_mb = []
         if num_microbatches <= 1:

src/twinkle/model/megatron/multi_lora_megatron.py

@@ -129,7 +129,7 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T
         with self.multi_adapter.adapter(adapter_name, disable_lora=disable_lora):
             return super().forward_only(inputs=inputs, **kwargs)
 
-    @remote_function(dispatch='slice_dp', collect='mean', sync=True)
+    @remote_function(dispatch='slice_dp', collect=collect_tensor_dict, sync=True)
     def forward_backward(self,
                          *,
                          inputs: Union[InputFeature, List[InputFeature], Trajectory, List[Trajectory]],

src/twinkle/model/multi_lora.py

@@ -205,7 +205,7 @@ def _linear_forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
 
                 lora_A_keys = self.lora_A.keys()
                 for active_adapter in self.active_adapters:
-                    if active_adapter not in lora_A_keys:
+                    if active_adapter not in lora_A_keys or self.disable_adapters:
                         continue
                     _lora = _self.find_lora(active_adapter)
                     target_modules = _lora.tenant_config.target_modules
@@ -238,7 +238,7 @@ def _embedding_forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
 
                 lora_embedding_A_keys = self.lora_embedding_A.keys()
                 for active_adapter in self.active_adapters:
-                    if active_adapter not in lora_embedding_A_keys:
+                    if active_adapter not in lora_embedding_A_keys or self.disable_adapters:
                         continue
                     _lora = self.find_lora(active_adapter)
                     target_modules = _lora.tenant_config.target_modules

src/twinkle/model/transformers/multi_lora_transformers.py

@@ -8,6 +8,7 @@
 from typing import Any, Callable, Dict, List, Literal, Optional, Type, Union
 
 from twinkle import DeviceMesh, remote_class, remote_function, template
+from twinkle.infra import collect_tensor_dict
 from twinkle.data_format import InputFeature, Trajectory
 from twinkle.hub import HubOperation
 from twinkle.loss import Loss
@@ -88,7 +89,7 @@ def unregister_mm_forward_hook(self, optimizer_group: OptimizerGroup):
     def _lazy_wrap_model(self):
         pass
 
-    @remote_function(dispatch='slice_dp', collect='mean')
+    @remote_function(dispatch='slice_dp', collect=collect_tensor_dict)
     def forward(self, *, inputs: Union[InputFeature, List[InputFeature], Trajectory, List[Trajectory]], **kwargs):
         self._check_adapter_valid(kwargs.get('adapter_name'))
         optimizer_config = self.optimizer_group[kwargs.get('adapter_name')]
@@ -104,7 +105,7 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], Trajectory,
         with self.multi_adapter.adapter(kwargs.get('adapter_name')):
             return super().forward(inputs=inputs, **kwargs)
 
-    @remote_function(dispatch='slice_dp', collect='flatten')
+    @remote_function(dispatch='slice_dp', collect=collect_tensor_dict)
     def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajectory]], **kwargs):
         adapter_name = kwargs.get('adapter_name')
         disable_lora = kwargs.get('disable_lora', False)
@@ -246,6 +247,7 @@ def set_grad_scaler(self, **kwargs):
         self._check_adapter_valid(kwargs.get('adapter_name'))
         super().set_grad_scaler(**kwargs)
 
+    @remote_function()
     def add_metric(self, metric_cls: Union[Metric, str], is_training: Optional[bool] = None, **kwargs):
         self._check_adapter_valid(kwargs.get('adapter_name'))
         super().add_metric(metric_cls, is_training, **kwargs)