[Fix] Convert lm_head and weights_proj to fp32

[bdubayah]

I was seeing the following error when running inference on the converted weights:

[rank0]: Traceback (most recent call last):
[rank0]: File "/workspace/TensorRT-Model-Optimizer/examples/deepseek/DeepSeek-V3.2-Exp/inference/generate.py", line 186, in <module>
[rank0]: main(args.ckpt_path, args.config, args.input_file, args.interactive, args.max_new_tokens, args.temperature)
[rank0]: File "/workspace/TensorRT-Model-Optimizer/examples/deepseek/DeepSeek-V3.2-Exp/inference/generate.py", line 119, in main
[rank0]: load_model(model, os.path.join(ckpt_path, f"model{rank}-mp{world_size}.safetensors"))
[rank0]: File "/opt/venv/lib/python3.12/site-packages/safetensors/torch.py", line 271, in load_model
[rank0]: raise RuntimeError(error)
[rank0]: RuntimeError: Error(s) in loading state_dict for Transformer:
[rank0]: Invalid key(s) in state_dict: "head.weight", "layers.0.attn.indexer.weights_proj.weight", "layers.1.attn.indexer.weights_proj.weight", "layers.10.attn.indexer.weights_proj.weight", "layers.11.attn.indexer.weights_proj.weight", "layers.12.attn.indexer.weights_proj.weight", "layers.13.attn.indexer.weights_proj.weight", "layers.14.attn.indexer.weights_proj.weight", "layers.15.attn.indexer.weights_proj.weight", "layers.16.attn.indexer.weights_proj.weight", "layers.17.attn.indexer.weights_proj.weight", "layers.18.attn.indexer.weights_proj.weight", "layers.19.attn.indexer.weights_proj.weight", "layers.2.attn.indexer.weights_proj.weight", "layers.20.attn.indexer.weights_proj.weight", "layers.21.attn.indexer.weights_proj.weight", "layers.22.attn.indexer.weights_proj.weight", "layers.23.attn.indexer.weights_proj.weight", "layers.24.attn.indexer.weights_proj.weight", "layers.25.attn.indexer.weights_proj.weight", "layers.26.attn.indexer.weights_proj.weight", "layers.27.attn.indexer.weights_proj.weight", "layers.28.attn.indexer.weights_proj.weight", "layers.29.attn.indexer.weights_proj.weight", "layers.3.attn.indexer.weights_proj.weight", "layers.30.attn.indexer.weights_proj.weight", "layers.31.attn.indexer.weights_proj.weight", "layers.32.attn.indexer.weights_proj.weight", "layers.33.attn.indexer.weights_proj.weight", "layers.34.attn.indexer.weights_proj.weight", "layers.35.attn.indexer.weights_proj.weight", "layers.36.attn.indexer.weights_proj.weight", "layers.37.attn.indexer.weights_proj.weight", "layers.38.attn.indexer.weights_proj.weight", "layers.39.attn.indexer.weights_proj.weight", "layers.4.attn.indexer.weights_proj.weight", "layers.40.attn.indexer.weights_proj.weight", "layers.41.attn.indexer.weights_proj.weight", "layers.42.attn.indexer.weights_proj.weight", "layers.43.attn.indexer.weights_proj.weight", "layers.44.attn.indexer.weights_proj.weight", "layers.45.attn.indexer.weights_proj.weight", "layers.46.attn.indexer.weights_proj.weight", "layers.47.attn.indexer.weights_proj.weight", "layers.48.attn.indexer.weights_proj.weight", "layers.49.attn.indexer.weights_proj.weight", "layers.5.attn.indexer.weights_proj.weight", "layers.50.attn.indexer.weights_proj.weight", "layers.51.attn.indexer.weights_proj.weight", "layers.52.attn.indexer.weights_proj.weight", "layers.53.attn.indexer.weights_proj.weight", "layers.54.attn.indexer.weights_proj.weight", "layers.55.attn.indexer.weights_proj.weight", "layers.56.attn.indexer.weights_proj.weight", "layers.57.attn.indexer.weights_proj.weight", "layers.58.attn.indexer.weights_proj.weight", "layers.59.attn.indexer.weights_proj.weight", "layers.6.attn.indexer.weights_proj.weight", "layers.60.attn.indexer.weights_proj.weight", "layers.7.attn.indexer.weights_proj.weight", "layers.8.attn.indexer.weights_proj.weight", "layers.9.attn.indexer.weights_proj.weight", mismatched dtypes or shape.

Start command:

torchrun --nproc-per-node 8 generate.py --ckpt-path path/to/weights --config ./config_671B_v3.2.json --interactive

After changing the dtypes to fp32 the error goes away. Conversion seems necessary based on the comments in model.py

    # weights_proj in the checkpoint is stored in bf16, while the parameters here are stored in fp32 for convenient.

[bdubayah]

cc @GeeeekExplorer