electron-rare/exo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: free broadcast weights before sharding to halve peak memory

Browse Source 
After model.load_weights(), both the broadcast_weights dict and the
model's parameter tree hold references to the same arrays. During
tensor_auto_parallel, the old full-size arrays can't be freed because
the dict still references them, causing ~2x peak memory.

Delete the dict before sharding so arrays are freed as each layer is
replaced with its sharded version.

Co-Authored-By: Claude Opus 4.6 <[email protected]>
This commit is contained in:
Alex Cheema
2026-02-17 10:05:43 -08:00
co-authored by Claude Opus 4.6
parent b594c2f58b
commit aab8aa38a9
1 changed files with 5 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/exo/worker/engines/mlx/utils_mlx.py
+5
View File
@@ -267,6 +267,11 @@ def shard_and_load(
)
model.load_weights(list(broadcast_weights.items()), strict=False)
# Free broadcast weight refs before sharding. The model's parameter tree is
# now the only holder of these arrays, so they'll be freed as each layer is
# replaced with its sharded version in tensor/pipeline_auto_parallel.
del broadcast_weights
tokenizer = get_tokenizer(model_path, shard_metadata)
logger.info(f"Group size: {group.size()}, group rank: {group.rank()}")