From aab8aa38a9ecd57dc0f1fa3a71b975b216b86cef Mon Sep 17 00:00:00 2001 From: Alex Cheema Date: Thu, 12 Feb 2026 13:14:12 -0800 Subject: [PATCH] fix: free broadcast weights before sharding to halve peak memory 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 --- src/exo/worker/engines/mlx/utils_mlx.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py index d299260e..e493dc6a 100644 --- a/src/exo/worker/engines/mlx/utils_mlx.py +++ b/src/exo/worker/engines/mlx/utils_mlx.py @@ -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()}")