[mutichung]

Summary

Remove extra nn.Embeddings quantizers that are not supported in ModelOpt v0.44.0.

Issue

Originally, edge was quantized with ModelOpt v0.47.0, while nano/super were quantized with v0.44.0. ModelOpt has nn.Embeddings support in v0.47.0, which leads to restoration failure when loading the checkpoint with v0.44.0. We decided to revert the ModelOpt version to v0.44.0 to make all Cosmos3 FP8 checkpoints consistent.

Example error message

raise ApplyModeError(f"Unmatched keys in quantizer state_dict: {unmatched_keys}")
modelopt.torch.opt.conversion.ApplyModeError: Unmatched keys in quantizer state_dict: {'embed_tokens.input_quantizer', 'action_proj_in.bias.input_quantizer', 'embed_tokens.weight_quantizer', 'action_proj_out.fc.output_quantizer', 'action_proj_out.bias.weight_quantizer', 'action_proj_in.bias.output_quantizer', 'action_proj_out.fc.input_quantizer', 'action_proj_in.fc.output_quantizer', 'action_proj_in.fc.input_quantizer', 'action_proj_in.bias.weight_quantizer', 'action_proj_out.fc.weight_quantizer', 'action_proj_in.fc.weight_quantizer', 'embed_tokens.output_quantizer', 'action_proj_out.bias.output_quantizer', 'action_proj_out.bias.input_quantizer'}

Fix

Additional quantizer metadata for nn.Embeddings are removed. Those layers remain bfloat16 regardless of ModelOpt version, so removing the layer metadata resolves the issue.