Fail to load state dict for Megatron Model after PTQ

[yikaizhu-baseten]

Before submitting an issue, please make sure it hasn't been already addressed by searching through the existing and past issues.

Describe the bug

https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/modelopt/torch/opt/plugins/megatron.py#L80
When I try to use NeMo to load model for finetuning after PTQ, I got errors in this line.
I have no idea what this line, but I have tested in locally and it fails

import torch
import pickle
x = torch.randint(0, 10, (2,))
y = pickle.loads(x.numpy().tobytes())

• ?

Steps/Code to reproduce bug

• ?

Expected behavior

Who can help?

• ?

System information

• Container used (if applicable): ?
• OS (e.g., Ubuntu 22.04, CentOS 7, Windows 10): ?
• CPU architecture (x86_64, aarch64): ?
• GPU name (e.g. H100, A100, L40S): ?
• GPU memory size: ?
• Number of GPUs: ?
• Library versions (if applicable):
  • Python: ?
  • ModelOpt version or commit hash: ?
  • CUDA: ?
  • PyTorch: ?
  • Transformers: ?
  • TensorRT-LLM: ?
  • ONNXRuntime: ?
  • TensorRT: ?
• Any other details that may help: ?

[realAsma]

Hi @yikaizhu-baseten could you please share your error log?

[yueshen2016]

Hi @yikaizhu-baseten , apart from the error log, could you also share the command you were running and the branch of the Megatron-Core and NeMo? And also the docker container you were using.

[yikaizhu-baseten]

I am using this docker nvcr.io/nvidia/nemo:25.09
First I use to get the nemo bf16 checkpoint

from nemo.collections import llm

if __name__ == "__main__":
  llm.import_ckpt(model=llm.Qwen3Model(llm.Qwen3Config8B()), source='hf://Qwen/Qwen3-8B')

then I use this where ptq.py is located in nemo/scripts/llm

torchrun \
	--nnodes 1 \
	--node_rank=0 \
	--master_addr=localhost\
	--nproc-per-node 4\
	--master_port=12346 \
	./ptq.py \
	--nemo_checkpoint=XXX \
	-cep=4 \
	-bs=16 \
	-sl=128 \
	-calib_size=16 \
	--algorithm=nvfp4 \
	--export_format=nemo \
	--export_path=YYY \
	--devices 4\
	-nodes 1

Finally, when i run gpt_train.py in nemo/scripts/llm I got this error

[yueshen2016]

@yikaizhu-baseten Could you share the error log? Thanks!