update

Author: jacky

xtuner/v1/datasets/sft_tokenize_fn/openai.py

@@ -37,8 +37,8 @@ def __init__(
         self.max_length = max_length
 
     def __call__(self, item: dict | list, **kwargs) -> DataItem | CacheItem:
-        if isinstance(item, dict) and "messages" in item:
-            item = item["messages"]
+        if isinstance(item, dict) and ("messages" in item or "dialogs" in item):
+            item = item.get("messages", item.get("dialogs"))
         messages = ChatMessages(messages=item)
         tokenized = messages.tokenize(self.tokenizer, self.chat_template)
 

xtuner/v1/ops/attn_imp.py

@@ -213,7 +213,7 @@ def flash_attention(q, k, v, window_size=(-1, -1), s_aux=None, **kwargs) -> torc
         if flash_attn_exception is not None:
             traceback.print_exception(flash_attn_exception)
             raise flash_attn_exception
-        attention_output = flash_attn_varlen_func(q, k, v, **kwargs)
+        attention_output = flash_attn_varlen_func(q, k, v, window_size=window_size, **kwargs)
     else:
         if flash_sink_attn_exception is not None:
             traceback.print_exception(flash_sink_attn_exception)

xtuner/v1/rl/base/worker.py

@@ -187,7 +187,7 @@ def _build_ref_model(
         if ref_model_fsdp_cfg is None:
             ref_model_fsdp_cfg = FSDPConfig(recompute_ratio=0, cpu_offload=False, requires_grad=False)
         model = model.fully_shard(ref_model_fsdp_cfg, float8_handler)
-        model.from_hf(hf_path=load_from)
+        model.from_hf(hf_path=load_from, strict=False)
         model.eval()
         if float8_handler is not None:
             # As the ref model is not updated, we only compute params' scales once

xtuner/v1/train/trainer.py

@@ -986,8 +986,10 @@ def _log_step(
         grad_norm: float,
     ):
         """Log the training step information."""
-        tgs = step_consumed_tokens / step_time
-        e2e_tgs = total_consumed_tokens / train_time
+        if step_consumed_tokens == 0:
+            logger.warning("step_consumed_tokens is 0 due to padding")
+        tgs = max(1, step_consumed_tokens / step_time)
+        e2e_tgs = max(1, total_consumed_tokens / train_time)
         lr = self._lr_scheduler.get_last_lr()[0]
 
         remaining_steps = self.total_step - self.cur_step