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 TransformersModel
+from twinkle.model import MegatronModel
 from twinkle.preprocessor import EmojiDPOProcessor
 from twinkle.processor import InputProcessor
 
@@ -157,15 +157,15 @@ def main():
         lora_dropout=0.05,
     )
 
-    policy_model = TransformersModel(
+    policy_model = MegatronModel(
         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('AdamW', lr=LEARNING_RATE, weight_decay=0.01)
-    policy_model.set_lr_scheduler('CosineAnnealingLR', T_max=MAX_STEPS, eta_min=LEARNING_RATE * 0.1)
+    policy_model.set_optimizer('default', lr=LEARNING_RATE, weight_decay=0.01)
+    policy_model.set_lr_scheduler('default', lr_decay_steps=MAX_STEPS)
 
     # Set up loss function and metrics
     loss_fn = DPOLoss(
@@ -191,13 +191,14 @@ 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, disable_lora=True)
+        ref_outputs = policy_model.forward_only(inputs=dpo_batch, micro_batch_size=2, disable_lora=True)
 
         # 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,
         )
 
         # Gradient clipping and optimizer step

src/twinkle/loss/dpo.py

@@ -326,10 +326,7 @@ def __call__(
             reference_chosen_logps = torch.zeros_like(policy_chosen_logps)
             reference_rejected_logps = torch.zeros_like(policy_rejected_logps)
         else:
-            raise ValueError(
-                "ref_logps or (ref_chosen_logps, ref_rejected_logps) must be provided "
-                "unless reference_free=True"
-            )
+            return LossOutput(loss=torch.tensor(0.0, device=chosen_logps.device), num_tokens=0)
 
         # Compute DPO loss
         dpo_loss = self._compute_dpo_loss(

src/twinkle/metric/dpo.py

@@ -80,27 +80,18 @@ def accumulate(self, inputs: Union[InputFeature, List[InputFeature]], outputs: M
             - kwargs['ref_outputs']: Optional reference model outputs with 'logps'
         """
         import torch
-
         logps = outputs.get('logps')
         if logps is None:
             return
 
         # Get labels from inputs
         if isinstance(inputs, list):
-            # Stack labels from list of inputs
-            labels_list = [torch.as_tensor(inp['labels']) for inp in inputs]
-            max_len = max(l.shape[0] for l in labels_list)
-            padded = []
-            for l in labels_list:
-                if l.shape[0] < max_len:
-                    pad = torch.full((max_len - l.shape[0],), self.ignore_index, dtype=l.dtype)
-                    l = torch.cat([pad, l])
-                padded.append(l)
-            labels = torch.stack(padded)
-        else:
-            labels = torch.as_tensor(inputs['labels'])
-            if labels.dim() == 1:
-                labels = labels.unsqueeze(0)
+            assert len(inputs) == 1
+            inputs = inputs[0]
+
+        labels = torch.as_tensor(inputs['labels'])
+        if labels.dim() == 1:
+            labels = labels.unsqueeze(0)
 
         # Ensure logps and labels have same device
         if logps.device != labels.device:
@@ -129,7 +120,6 @@ def accumulate(self, inputs: Union[InputFeature, List[InputFeature]], outputs: M
             ref_logps = ref_outputs.get('logps')
             if ref_logps is not None:
                 # Align ref_logps to match labels shape (handles different seq lengths)
-                # breakpoint()
                 ref_logps = self._align_logps(
                     ref_logps, labels.shape, labels.device, logps.dtype
                 )

src/twinkle/model/megatron/megatron.py

@@ -447,7 +447,13 @@ def post_loss_function(output_tensor, inputs, logps):
         def forward_step_func(data_iterator, model):
             batch = next(data_iterator)
             labels = batch.pop('labels', None)
-            output_tensor = model(**batch)
+            # Handle disable_lora for base model inference (e.g., reference in DPO)
+            unwrapped_model = self.strategy.unwrap_model([model])[0]
+            if disable_lora and isinstance(unwrapped_model, PeftModel):
+                with unwrapped_model.disable_adapter():
+                    output_tensor = model(**batch)
+            else:
+                output_tensor = model(**batch)
             batch['labels'] = labels
             logps = None
             if labels is not None and mpu.is_pipeline_last_stage():
@@ -475,34 +481,17 @@ def forward_step_func(data_iterator, model):
 
         self._accumulate_metric(optimizer_config, is_training=not forward_only)
 
-        # Handle disable_lora for base model inference (e.g., reference in DPO)
-        def _set_disable_adapters(model, value: bool):
-            model = self.strategy.unwrap_model(model)
-            if isinstance(model, list):
-                for m in model:
-                    if isinstance(m, PeftModel):
-                        m.disable_adapters = value
-            elif isinstance(model, PeftModel):
-                model.disable_adapters = value
-
-        if disable_lora:
-            _set_disable_adapters(self.model, True)
-
-        try:
-            # Run forward-backward with Megatron's scheduler
-            # Megatron handles all communication internally using proper process groups
-            losses = forward_backward_func(
-                forward_step_func=forward_step_func,
-                data_iterator=data_iter,
-                model=self.model,
-                num_microbatches=len(inputs),
-                seq_length=seq_length,
-                micro_batch_size=micro_batch_size,
-                forward_only=forward_only,
-            )
-        finally:
-            if disable_lora:
-                _set_disable_adapters(self.model, False)
+        # Run forward-backward with Megatron's scheduler
+        # Megatron handles all communication internally using proper process groups
+        losses = forward_backward_func(
+            forward_step_func=forward_step_func,
+            data_iterator=data_iter,
+            model=self.model,
+            num_microbatches=len(inputs),
+            seq_length=seq_length,
+            micro_batch_size=micro_batch_size,
+            forward_only=forward_only,
+        )
 
         # Extract loss from results (only last PP stage returns non-empty)
         loss = torch.tensor(0.0).to(Platform.get_local_device())
@@ -559,9 +548,11 @@ def _set_disable_adapters(model, value: bool):
         if forward_only:
             optimizer_config.eval_status.inputs = inputs
             optimizer_config.eval_status.outputs = ModelOutput(logits=logits, loss=loss, logps=logps)
+            optimizer_config.eval_status.forward_kwargs = kwargs
         else:
             optimizer_config.train_status.inputs = inputs
             optimizer_config.train_status.outputs = ModelOutput(logits=logits, loss=loss, logps=logps)
+            optimizer_config.train_status.forward_kwargs = kwargs
         return ModelOutput(logits=logits, loss=loss, logps=logps)
 
     @remote_function(dispatch='all')
@@ -692,6 +683,7 @@ def set_loss(self, loss_cls: Union[Loss, Type[Loss], str, Callable[[InputFeature
         optimizer_config = self.optimizer_group[adapter_name]
         optimizer_config.loss_instance = construct_class(loss_cls, Loss, twinkle.loss, **kwargs)
 
+    @remote_function()
     def add_metric(self, metric_cls: Union[Metric, str], is_training: Optional[bool] = None, **kwargs):
         """Add an eval metric
 
@@ -773,16 +765,16 @@ def _create_megatron_optimizer(self, **kwargs):
         opt_config = OptimizerConfig(
             optimizer='adam',
             lr=lr,
-            min_lr=kwargs.get('min_lr', 0.0),
-            weight_decay=kwargs.get('weight_decay', 0.01),
-            adam_beta1=kwargs.get('adam_beta1', 0.9),
-            adam_beta2=kwargs.get('adam_beta2', 0.999),
-            adam_eps=kwargs.get('adam_eps', 1e-8),
-            clip_grad=kwargs.get('clip_grad', 1.0),
-            bf16=kwargs.get('bf16', True),
+            min_lr=kwargs.pop('min_lr', 0.0),
+            weight_decay=kwargs.pop('weight_decay', 0.01),
+            adam_beta1=kwargs.pop('adam_beta1', 0.9),
+            adam_beta2=kwargs.pop('adam_beta2', 0.999),
+            adam_eps=kwargs.pop('adam_eps', 1e-8),
+            clip_grad=kwargs.pop('clip_grad', 1.0),
+            bf16=kwargs.pop('bf16', True),
             use_distributed_optimizer=use_distributed_optimizer,
-            overlap_param_gather=kwargs.get('overlap_param_gather', False),
-            log_num_zeros_in_grad=kwargs.get('log_num_zeros_in_grad', False),
+            overlap_param_gather=kwargs.pop('overlap_param_gather', False),
+            log_num_zeros_in_grad=kwargs.pop('log_num_zeros_in_grad', False),
             **kwargs,
         )
 

src/twinkle/model/megatron/multi_lora_megatron.py

@@ -264,6 +264,7 @@ def set_processor(self, processor_cls: Union[Type[InputProcessor], str, Callable
         self._check_adapter_valid(kwargs.get('adapter_name'))
         super().set_processor(processor_cls, **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)

src/twinkle/model/transformers/transformers.py

@@ -188,7 +188,6 @@ def __init__(
         }
         self.optimizer_group[_default_adapter_name].adapter_name = _default_adapter_name
         self.active_group = _default_adapter_name
-        # breakpoint()
 
     def _decide_strategy(self, strategy: Literal['accelerate', 'native_fsdp']):
         self._expert_parallel_config = self._fsdp_config.pop('expert_parallel', None)