Skip to content

Categorize optimizer step and construction frames as OPTIMIZER#20

Closed
Chenguang-Zhu wants to merge 1 commit into
facebookresearch:mainfrom
Chenguang-Zhu:export-D103992613
Closed

Categorize optimizer step and construction frames as OPTIMIZER#20
Chenguang-Zhu wants to merge 1 commit into
facebookresearch:mainfrom
Chenguang-Zhu:export-D103992613

Conversation

@Chenguang-Zhu

Copy link
Copy Markdown
Member

Summary:
Expand the OPTIMIZER name set and add a filename-based OPTIMIZER rule
to capture vanilla torch.optim, TorchRec keyed optimizer, Distributed
Shampoo, and additional in-house optimizer implementations.

Public name additions in _OPTIMIZER_VANILLA_NAMES:

  • step, _per_group_step_impl
  • create_optimizer, _create_optimizer, create_keyed_optimizer
  • _instantiate_shampoo_preconditioner_list,
    _create_kronecker_factors_state,
    _create_kronecker_factors_state_for_block,
    _create_base_kronecker_factors

Public filename-substring rule (_OPTIMIZER_FILENAME_SUBSTRINGS):

  • torch/optim/optimizer.py
  • torchrec/optim/keyed.py

In-house optimizer rules are injected via the
mosaic.libmosaic.utils.internal_optimizer_rules extension hook so the
OSS-shipping data_utils.py carries no internal-only path strings or
symbol names.

Differential Revision: D103992613

Summary:
Expand the OPTIMIZER name set and add a filename-based OPTIMIZER rule
to capture vanilla `torch.optim`, TorchRec keyed optimizer, Distributed
Shampoo, and additional in-house optimizer implementations.

Public name additions in `_OPTIMIZER_VANILLA_NAMES`:
- step, _per_group_step_impl
- create_optimizer, _create_optimizer, create_keyed_optimizer
- _instantiate_shampoo_preconditioner_list,
  _create_kronecker_factors_state,
  _create_kronecker_factors_state_for_block,
  _create_base_kronecker_factors

Public filename-substring rule (`_OPTIMIZER_FILENAME_SUBSTRINGS`):
- torch/optim/optimizer.py
- torchrec/optim/keyed.py

In-house optimizer rules are injected via the
`mosaic.libmosaic.utils.internal_optimizer_rules` extension hook so the
OSS-shipping `data_utils.py` carries no internal-only path strings or
symbol names.

Differential Revision: D103992613
@meta-cla meta-cla Bot added the cla signed label May 15, 2026
@meta-codesync

meta-codesync Bot commented May 15, 2026

Copy link
Copy Markdown

@Chenguang-Zhu has exported this pull request. If you are a Meta employee, you can view the originating Diff in D103992613.

@meta-codesync

meta-codesync Bot commented May 15, 2026

Copy link
Copy Markdown

This pull request has been merged in 1b2bfb3.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant