SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

Fix load balancing loss func for mixtral (#28256)

Browse Source 
* Correct the implementation of auxiliary loss of mixtrtal

* correct the implementation of auxiliary loss of mixtrtal

* Implement a simpler calculation method

---------

Co-authored-by: zhangliangxu3 <zhangliangxu3@jd.com>
This commit is contained in:
liangxuZhang
2024-01-11 23:16:12 +08:00
committed by GitHub
parent 5d4d62d0a2
commit e768616afa
2 changed files with 2 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tests/models/mixtral/test_modeling_mixtral.py
View File

@@ -474,7 +474,7 @@ class MixtralModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMi
model.eval()
result = model(input_ids, attention_mask=attention_mask)
self.assertEqual(result.router_logits[0].shape, (91, config.num_local_experts))
torch.testing.assert_close(result.aux_loss.cpu(), torch.tensor(8, dtype=torch.float32))
torch.testing.assert_close(result.aux_loss.cpu(), torch.tensor(2, dtype=torch.float32), rtol=1e-2, atol=1e-2)
@require_torch