Lower complexity of get_balanced_memory by adding a set

[SamuelBarryCS]

What

• Fixes Optimize the time complexity of get_balanced_memory #3768 , i.e. reduces the complexity of checking whether modules are leaves by switching from repeated list membership checks (O(N)) to a set-based approach (O(1) average).
• This optimization lowers the overhead of loading large models, especially MoEs!

How to review

• Read diff

Tests performed

• Functionality/logic unchanged -> existing tests continue to pass.

Benchmark results

• On a representative setup of 3 models (one small, large and very large MoE) , runtime scales much better with the set-based approach.
• TLDR: ~10–50× faster membership handling for large module lists.

[SamuelBarryCS]

cc @Jintao-Huang @S1ro1 for review/ stamp :)
Thanks!

[S1ro1]

Can you please clean up before merge? will stamp after. LGTM rn.

[SamuelBarryCS]

Can you please clean up before merge? will stamp after. LGTM rn.

Sure thing @S1ro1 - done now!
Thanks!

[SunMarc]

Thanks !

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.