Skip to content

Optimize the time complexity of get_balanced_memory #3768

New issue
New issue
Closed
#3776
Closed
Optimize the time complexity of get_balanced_memory#3768
#3776
@Jintao-Huang

Description

@Jintao-Huang
Jintao-Huang
opened on Sep 9, 2025

When loading very large MoE models it stalls here; I found this part has O(n^2) complexity. If you convert it to a set, it can be reduced to O(n).

accelerate/src/accelerate/utils/modeling.py

Lines 1043 to 1045 in 8b49352

leaves = get_module_leaves(module_sizes)
module_sizes = {n: v for n, v in module_sizes.items() if n not in leaves}
# Once removed, leaves are the final modules.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions