fix loss computation bug

Author: meichangsu1

cookbook/legacy/sft/single_controller_sp_packing.py

@@ -8,6 +8,7 @@
 from twinkle.dataset import PackingDataset, DatasetMeta
 from twinkle.model import TransformersModel
 from twinkle.preprocessor import SelfCognitionProcessor
+from twinkle.processor import InputProcessor
 
 logger = get_logger()
 MODEL_ID = 'ms://Qwen/Qwen2.5-7B-Instruct'
@@ -82,9 +83,10 @@ def train():
         strategy="native_fsdp",
         remote_group="default",
     )
-
     lora_config = LoraConfig(target_modules="all-linear")
     model.add_adapter_to_model("default", lora_config, gradient_accumulation_steps=1)
+    model.set_processor(InputProcessor, padding_free=True, adapter_name="default")
+    model.set_loss("CrossEntropyLoss", reduction="mean", adapter_name="default")
     model.set_optimizer("AdamW", lr=1e-4, adapter_name="default")
 
     loss_metric = 99.0

cookbook/legacy/single_controller_sp.py

@@ -46,27 +46,13 @@ def create_dataset(data_slice=None):
         "Template",
         model_id=MODEL_ID,
         truncation_strategy="left",
-        max_length=64,
+        max_length=256,
     )
     dataset.map(SelfCognitionProcessor("twinkle模型", "twinkle团队"))
     dataset.encode(batched=True)
     return dataset
 
 
-def eval(model: TransformersModel):
-    dataloader = DataLoader(
-        dataset=partial(create_dataset, data_slice=range(20)),
-        batch_size=4,
-        drop_last=True,
-        device_mesh=device_mesh,
-        remote_group="default",
-    )
-    for step, batch in enumerate(dataloader):
-        model.forward_only(inputs=batch, adapter_name="default")
-        model.calculate_loss(adapter_name="default")
-    metrics = model.calculate_metric(is_training=False, adapter_name="default")
-    return metrics()
-
 
 def train():
     dataloader = DataLoader(
@@ -87,21 +73,20 @@ def train():
     model.add_adapter_to_model("default", lora_config, gradient_accumulation_steps=1)
     model.set_optimizer("AdamW", lr=1e-4, adapter_name="default")
 
-    loss_metric = 99.0
     for step, batch in enumerate(dataloader):
-        if isinstance(batch, list) and len(batch) == 0:
-            continue
-        output = model.forward_backward(inputs=batch, adapter_name="default")
-        loss_value = output() if callable(output) else output
-        logger.info(f"step {step}, loss: {loss_value}")
-        model.clip_grad_and_step(adapter_name="default")
-        if step % 50 == 0 and step > 0:
-            metrics = eval(model)
-            logger.info(f"Current is step {step} of {len(dataloader)}, metric: {metrics}")
-            metrics["step"] = step
-            if loss_metric > metrics["loss"]:
-                model.save(f"checkpoint-{step}")
-                loss_metric = metrics["loss"]
+        model.forward_backward(inputs=batch, adapter_name='default')
+        model.clip_grad_and_step(adapter_name='default')
+        if step % 1 == 0:
+            metric = model.calculate_metric(is_training=True, adapter_name='default')
+            _metrics = {}
+            for key, value in metric.items():
+                try:
+                    value = float(value)
+                    _metrics[key] = value
+                except:
+                    pass
+            logger.info(f'Current is step {step} of {len(dataloader)}, metric: {metric}')
+    model.save(f'last-checkpoint', interval=1)
 
 
 if __name__ == "__main__":

src/twinkle/model/transformers/strategy/sequence_parallel.py

@@ -1,5 +1,6 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
 import math
+import os
 from functools import partial
 from types import SimpleNamespace
 from typing import Any, Dict, Optional, Tuple, Union
@@ -1004,7 +1005,11 @@ def reduce_loss(self, loss: torch.Tensor, labels: Optional[torch.Tensor], ignore
             local_sum = loss
             global_sum = local_sum.detach().clone()
             dist.all_reduce(global_sum, group=sequence_parallel._sp_group)
-            return global_sum + (local_sum - local_sum.detach())
+            out = global_sum + (local_sum - local_sum.detach())
+            if sequence_parallel.world_size > 1:
+                out_metric = out.detach() / sequence_parallel.world_size
+                return out_metric + (out - out.detach())
+            return out
         # Default to mean reduction.
         local_sum = loss * num_valid_tokens
         global_sum = local_sum.detach().clone()
@@ -1013,7 +1018,11 @@ def reduce_loss(self, loss: torch.Tensor, labels: Optional[torch.Tensor], ignore
         dist.all_reduce(global_tokens, group=sequence_parallel._sp_group)
         if global_tokens.item() == 0:
             return loss
-        return (global_sum + (local_sum - local_sum.detach())) / global_tokens
+        out = (global_sum + (local_sum - local_sum.detach())) / global_tokens
+        if sequence_parallel.world_size > 1:
+            out_metric = out.detach() / sequence_parallel.world_size
+            return out_metric + (out - out.detach())
+        return out
 
     def wrap_model(self, model, optimizer=None):
         self.initialize()

src/twinkle/model/transformers/transformers.py

@@ -1,5 +1,6 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
 import contextlib
+import os
 import json
 import os
 import re