fix: 修复部分已知问题

khaosz/inference/core.py

@@ -191,7 +191,7 @@ class EmbeddingEncoderCore:
         sentence_embs: List[Tensor] = []
         for i in range(len(batch_ids)):
             indices = [idx for idx, orig_idx in enumerate(fragment_origin_idx) if orig_idx == i]
-            if indices is not None:
+            if indices:
                 sum_frags = torch.sum(fragment_embs[indices, :, :], dim=1)      # [frags, hidden_size]
                 length = torch.sum(seq_mask[indices, :], dim=1).unsqueeze(1)    # [frags, 1]
                 emb = torch.sum(sum_frags / length, dim=0)                      # [frags, hidden_size]
@@ -228,11 +228,11 @@ class KVCacheManager:
         self._initialize()
     
     def _initialize(self):
-        k_cache = torch.zeros(
+        k_cache = torch.empty(
             (self.batch_size, self.max_len, self.num_layers, self.num_heads, self.head_dim),
             device=self.device, dtype=self.dtype
         )
-        v_cache = torch.zeros(
+        v_cache = torch.empty(
             (self.batch_size, self.max_len, self.num_layers, self.num_heads, self.head_dim),
             device=self.device, dtype=self.dtype
         )

khaosz/model/module.py

@@ -93,7 +93,7 @@ class RotaryEmbedding(nn.Module):
         seq_len = x.size(1)
         
         if self.max_len_cached < seq_len + start_pos:
-            self._set_rotary_buffer(seq_len)
+            self._set_rotary_buffer(seq_len + start_pos)
         
         cos = self.cos_cached[start_pos : start_pos + seq_len]
         sin = self.sin_cached[start_pos : start_pos + seq_len]
@@ -237,6 +237,7 @@ class MLA(nn.Module):
         use_gated_attention: bool,
         layer_id: int
     ):
+        super().__init__()
         self.dim = dim
         self.n_heads = n_heads
         self.n_kv_heads = n_kv_heads

khaosz/parallel/setup.py

@@ -82,9 +82,12 @@ def only_on_rank(rank, sync=False):
         @wraps(func)
         def wrapper(*args, **kwargs):
             if get_rank() == rank:
-                return func(*args, **kwargs)
-            if sync:
+                ret_args = func(*args, **kwargs)
+
+            if sync and dist.is_available() and dist.is_initialized():
                 dist.barrier()
+
+            return ret_args
         
         return wrapper
     

khaosz/trainer/train_callback.py

@@ -74,19 +74,16 @@ class SchedulerCallback(TrainCallback):
     Scheduler callback for trainer.
     """
     def __init__(self):
-        self.scheduler: LRScheduler = None
+        pass
     
     def on_train_begin(self, context: TrainContext):
         for group in context.optimizer.param_groups:
             if "initial_lr" not in group:
                 group["initial_lr"] = group["lr"] 
-        
-        self.scheduler = context.scheduler
     
     def on_batch_end(self, context: TrainContext):
-        _ = context
-        if self.scheduler:
-            self.scheduler.step()
+        if context.scheduler:
+            context.scheduler.step()
 
 
 class CheckpointCallback(TrainCallback):

khaosz/trainer/train_context.py

@@ -87,7 +87,7 @@ class TrainContextBuilder:
         return self
     
     def with_strategy(self) -> Self:
-        self._context.strategy = StrategyFactory.load(
+        self._context.strategy = StrategyFactory.create(
             model=self._context.model,
             train_type=self.config.strategy,
             device=get_current_device(),