Skip to content

Conversation

lisjin
Copy link
Contributor

@lisjin lisjin commented Sep 30, 2025

Fix compatibility with open source QAT training/finetuning scripts.

Details: FSDP2 state_dict and load_state_dict methods will raise errors if param_group contains a quantizer object. (This is necessary for quantizing embeddings differently from other params.) I replaced the param_group["quantizer"] key with a group_quantizer_map mapping between regularized_param_groups indices and quantizers in QuantOptimizer.

@lisjin lisjin requested a review from metascroy September 30, 2025 19:34
Copy link

pytorch-bot bot commented Sep 30, 2025

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/3104

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit abc60fd with merge base d407246 (image):
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@meta-cla meta-cla bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Sep 30, 2025
@lisjin lisjin added the topic: bug fix Use this tag for PRs that fix bugs label Oct 1, 2025
def get_optim_kwargs(base_optimizer, embedding=True, quant_cls=UnifTorchaoQuantizer):
optim_kwargs = {}
if embedding:
group_idx = len(base_optimizer.param_groups) - 2
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why this? It seems brittle?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I agree—I just changed it to be more generalizable

@lisjin lisjin force-pushed the lvj/group-quantizer branch from 2fbb4a0 to abc60fd Compare October 2, 2025 18:30
@lisjin lisjin requested a review from metascroy October 3, 2025 15:15
@lisjin lisjin merged commit f5bb94a into main Oct 3, 2025
18 checks passed
@lisjin lisjin deleted the lvj/group-quantizer branch October 3, 2025 19:09
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. topic: bug fix Use this tag for PRs that fix bugs

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants