fix calculate_shard_storages to handle optimizer correctly

[tiankongdeguiji]

The parameter no longer possesses the optimizer_class attribute; instead, it has been updated to optimizer_classes. However, the current implementation of the EmbeddingStorageEstimator still relies on the outdated optimizer_class attribute to determine whether the parameter includes an optimizer. As a result, the EmbeddingStorageEstimator fails to estimate the storage requirements for the optimizer, leading to CUDA out-of-memory errors during training.

[tiankongdeguiji]

Hi, @henrylhtsang @joshuadeng @PaulZhang12 @TroyGarden can you take a look?

[tiankongdeguiji]

hi, @sarckk @dstaay-fb can you take a look?

[facebook-github-bot]

@sarckk has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[sarckk]

@tiankongdeguiji thanks for the fix! looks good to me. I'll help you land the changes (EDIT: probably after 1st Jan)

[tiankongdeguiji]

@tiankongdeguiji thanks for the fix! looks good to me. I'll help you land the changes

thx!

[tiankongdeguiji]

hi, @sarckk can you help me land the changes? Are there any existing issues I should address?

[sarckk]

hi, @sarckk can you help me land the changes? Are there any existing issues I should address?

Hi sorry for the delay. This increases HBM and DDR estimates across the board so I'm running extra checks to avoid any regressions internally. I will land it by this week

[tiankongdeguiji]

hi, @sarckk can you help me land the changes? Are there any existing issues I should address?

Hi sorry for the delay. This increases HBM and DDR estimates across the board so I'm running extra checks to avoid any regressions internally. I will land it by this week

thx!

[sarckk]

@tiankongdeguiji thanks for the patience. your diff has been landed