Merge branch 'master' into rearrange_ops

microsoft · May 23, 2024 · 3034ce8 · 3034ce8
2 parents b8faade + 3a3a6db
commit 3034ce8
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/.github/workflows/nv-ds-chat.yml b/.github/workflows/nv-ds-chat.yml
@@ -14,6 +14,7 @@ on:
  paths:
  - "deepspeed/runtime/zero/stage_1_and_2.py"
  - "deepspeed/runtime/zero/stage3.py"
+ - "deepspeed/runtime/hybrid_engine.py"
 
 concurrency:
  group: ${{ github.workflow }}-${{ github.ref }}

diff --git a/deepspeed/sequence/layer.py b/deepspeed/sequence/layer.py
@@ -81,7 +81,7 @@ def __init__(
  self.scatter_idx = scatter_idx
  self.gather_idx = gather_idx
 
- def forward(self, query: Tensor, key: Tensor, value: Tensor, *args: Any) -> Tensor:
+ def forward(self, query: Tensor, key: Tensor, value: Tensor, *args: Any, **kwargs) -> Tensor:
  """ forward
 
  Arguments:
@@ -101,7 +101,7 @@ def forward(self, query: Tensor, key: Tensor, value: Tensor, *args: Any) -> Tens
  value_layer = _SeqAllToAll.apply(self.spg, value, self.scatter_idx, self.gather_idx)
 
  #out shape : e.g., [s:h/p:]
- context_layer = self.local_attn(query_layer, key_layer, value_layer, *args)
+ context_layer = self.local_attn(query_layer, key_layer, value_layer, *args, **kwargs)
 
  output = _SeqAllToAll.apply(self.spg, context_layer, self.gather_idx, self.scatter_idx)