RYDE-WORK/ktransformers
1
0
Fork 0
mirror of https://github.com/RYDE-WORK/ktransformers.git synced 2026-01-19 12:43:16 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Merge pull request #684 from KMSorSMS/main

Browse Source 
fix dockerfile in devcontainer and fix expert torch
This commit is contained in:
ZiWei Yuan 2025-02-26 15:06:51 +08:00 committed by GitHub
commit e7ebb26370
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 3 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
.devcontainer/Dockerfile
View File

@ -10,7 +10,6 @@ apt update -y && apt install -y --no-install-recommends \
g++ \
cmake &&
rm -rf /var/lib/apt/lists/* &&
cd ktransformers &&
pip install ninja pyproject numpy cpufeature &&
pip install flash-attn &&
cp /usr/lib/x86_64-linux-gnu/libstdc++.so.6 /opt/conda/lib/

6
ktransformers/operators/experts.py
View File

@ -459,9 +459,9 @@ class KExpertsTorch(KExpertsBase):
self.up[i] = w["up"][i, ...].to(device=device, dtype=self.dtype)
self.down[i] = w["down"][i, ...].to(device=device, dtype=self.dtype)
self.up = torch.cat(self.up, dim=0)
self.gate = torch.cat(self.gate, dim=0)
self.down = torch.cat(self.down, dim=0)
self.up = torch.stack(self.up, dim=0)
self.gate = torch.stack(self.gate, dim=0)
self.down = torch.stack(self.down, dim=0)
return
def unload(self):