saymrwulf/transformers
1
0
Fork 0
mirror of https://github.com/saymrwulf/transformers.git synced 2026-05-14 20:58:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Update test_tp.py

Browse source
This commit is contained in:
Cyril Vallez 2025-01-31 22:41:02 +01:00
parent 8c419c6fa9
commit a3a55d06df
No known key found for this signature in database
1 changed files with 2 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
tests/tp/test_tp.py
View file

@ -84,8 +84,8 @@ class TestTensorParallel(TestCasePlus):
expected_model_memory = 16
overhead_factor = 1.2
# # Check that we do not use more than the expected sharded size during initialization
if not torch.cuda.max_memory_allocated(device) / 1024**3 < (expected_model_memory / world_size) * overhead_factor:
# Check that we do not use more than the expected sharded size during initialization
if torch.cuda.max_memory_allocated(device) / 1024**3 > (expected_model_memory / world_size) * overhead_factor:
raise ValueError("Loading the model used more than the expected fraction of model size per device")
torch.distributed.barrier()