Skip to content

Fix loading of aliased weights - #85

Merged
suyoggupta merged 2 commits into
feat/ad-2025-07-07from
user/sg/aliasing-fix
Jul 14, 2025
Merged

Fix loading of aliased weights#85
suyoggupta merged 2 commits into
feat/ad-2025-07-07from
user/sg/aliasing-fix

Conversation

@suyoggupta

Copy link
Copy Markdown
  • Adds a test for small Qwen model that has aliased lm_head.weight and embed.weight
  • Removes a conservative assert that required aliased weights to be present in the ckpt. In cases when ckpt is loaded in parts, this assert would misfire.

@suyoggupta
suyoggupta requested review from Fridah-nv and Copilot July 14, 2025 19:05

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

Pull Request Overview

This PR updates test configurations to include the Qwen2.5-3B-Instruct model and relaxes loading aliases when checkpoint parts are missing.

  • Added a new small-model entry for Qwen/Qwen2.5-3B-Instruct in two test files
  • Removed a strict assert in aliasing_load_pre_hook to allow missing aliased weights without failure

Reviewed Changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.

File Description
tests/unittest/_torch/auto_deploy/unit/singlegpu/test_ad_build_small_single.py Added a get_small_model_config entry for Qwen2.5-3B-Instruct
tests/unittest/_torch/auto_deploy/_utils_test/_model_test_utils.py Registered Qwen2.5-3B-Instruct in the small-model lookup table
tensorrt_llm/_torch/auto_deploy/transformations/export.py Removed assert on missing alias load and wrapped aliasing in an if
Comments suppressed due to low confidence (1)

Comment thread tensorrt_llm/_torch/auto_deploy/transformations/export.py

@Fridah-nv Fridah-nv left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

LGTM, thank you!

Signed-off-by: Suyog Gupta <suyogg@nvidia.com>
@suyoggupta
suyoggupta enabled auto-merge (squash) July 14, 2025 19:16
@suyoggupta
suyoggupta merged commit 8b68352 into feat/ad-2025-07-07 Jul 14, 2025
1 check passed
lucaslie pushed a commit that referenced this pull request Jul 18, 2025
* remove assert, add qwen small to tests

* lint

Signed-off-by: Suyog Gupta <suyogg@nvidia.com>

---------

Signed-off-by: Suyog Gupta <suyogg@nvidia.com>
lucaslie pushed a commit that referenced this pull request Jul 21, 2025
* remove assert, add qwen small to tests

* lint

Signed-off-by: Suyog Gupta <suyogg@nvidia.com>

---------

Signed-off-by: Suyog Gupta <suyogg@nvidia.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants