core.post_training.modelopt.gpt.state_dict_hooks#

Module Contents#

Functions#

mcore_gpt_load_te_state_dict_pre_hook

Register a pre-hook to fix the state_dict key difference of.

Data#

logger

API#

core.post_training.modelopt.gpt.state_dict_hooks.logger#

‘getLogger(…)’

core.post_training.modelopt.gpt.state_dict_hooks.mcore_gpt_load_te_state_dict_pre_hook(
state_dict,
prefix,
local_metadata,
strict,
missing_keys,
unexpected_keys,
error_msgs,
)#

Register a pre-hook to fix the state_dict key difference of.

This prehook is used when trying to load the megatron/core GPTModel that uses a fused Transformer-Engine ParallelLinear into the variant that uses native ParallelLinear and Transformer-Engine Norm (effectively to restore the fusion). Only this particular spec supports post-training quantization and TensorRT-LLM config export through nvidia-modelopt package.

Parameters:

  • state_dict – state dictionary

  • prefix – module name prefix

  • local_metadata – local metatdata

  • strict – whether is in strict mode

  • missing_keys – missing state dict keys

  • unexpected_keys – unexpected state dict keys

  • error_msgs – error messages