Skip to content

Commit 895d451

Browse files
JiaxinDJiaxinD
committed
fix(checkpoints): preserve validation under python -O
Replace checkpoint-boundary assertions with explicit ValueError guards so malformed external data is still rejected when Python optimization strips assertions. Add an optimized-Python regression and enforce the audited internal-only assertion boundary. Refs: #1052 Signed-off-by: JiaxinD <djx2048@gmail.com>
1 parent 7aa0b21 commit 895d451

60 files changed

Lines changed: 308 additions & 176 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎python/tensorrt_model_connect/families/albert/plugin.py‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -68,20 +68,20 @@ def load_weights(
6868

6969
# Word embedding: [vocab, embedding_size]
7070
embedding = _load_tensor(readers, "albert.embeddings.word_embeddings.weight")
71-
assert embedding.shape == (vocab, embedding_size), (
72-
f"Embedding shape {embedding.shape} != ({vocab}, {embedding_size})")
71+
if embedding.shape != (vocab, embedding_size):
72+
raise ValueError(f'Embedding shape {embedding.shape} != ({vocab}, {embedding_size})')
7373
weights["embedding"] = embedding.astype(np.float32)
7474

7575
# Position embedding: [max_pos, embedding_size]
7676
pos_embed = _load_tensor(readers, "albert.embeddings.position_embeddings.weight")
77-
assert pos_embed.shape == (max_pos, embedding_size), (
78-
f"Position embedding shape {pos_embed.shape} != ({max_pos}, {embedding_size})")
77+
if pos_embed.shape != (max_pos, embedding_size):
78+
raise ValueError(f'Position embedding shape {pos_embed.shape} != ({max_pos}, {embedding_size})')
7979
weights["position_embedding"] = pos_embed.astype(np.float32)
8080

8181
# Token type embedding: [type_vocab_size, embedding_size]
8282
tt_embed = _load_tensor(readers, "albert.embeddings.token_type_embeddings.weight")
83-
assert tt_embed.shape == (type_vocab_size, embedding_size), (
84-
f"Token type embedding shape {tt_embed.shape} != ({type_vocab_size}, {embedding_size})")
83+
if tt_embed.shape != (type_vocab_size, embedding_size):
84+
raise ValueError(f'Token type embedding shape {tt_embed.shape} != ({type_vocab_size}, {embedding_size})')
8585
weights["token_type_embedding"] = tt_embed.astype(np.float32)
8686

8787
# Embedding LayerNorm (over embedding_size dim)

‎python/tensorrt_model_connect/families/bert/weights/__init__.py‎

Lines changed: 6 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -168,27 +168,22 @@ def load_bert_weights(model_dir: str | Path, config: ModelConfig) -> WeightDict:
168168
weights = WeightDict()
169169

170170
embedding = _load_tensor(readers, _prefixed(root, "embeddings.word_embeddings.weight"))
171-
assert embedding.shape == (config.vocab_size, hidden), (
172-
f"Embedding shape {embedding.shape} != ({config.vocab_size}, {hidden})"
173-
)
171+
if embedding.shape != (config.vocab_size, hidden):
172+
raise ValueError(f'Embedding shape {embedding.shape} != ({config.vocab_size}, {hidden})')
174173
weights["embedding"] = embedding.astype(np.float32)
175174

176175
position_embedding = _load_tensor(
177176
readers, _prefixed(root, "embeddings.position_embeddings.weight")
178177
)
179-
assert position_embedding.shape == (config.max_position_embeddings, hidden), (
180-
f"Position embedding shape {position_embedding.shape} != "
181-
f"({config.max_position_embeddings}, {hidden})"
182-
)
178+
if position_embedding.shape != (config.max_position_embeddings, hidden):
179+
raise ValueError(f'Position embedding shape {position_embedding.shape} != ({config.max_position_embeddings}, {hidden})')
183180
weights["position_embedding"] = position_embedding.astype(np.float32)
184181

185182
token_type_key = _prefixed(root, "embeddings.token_type_embeddings.weight")
186183
if _has_tensor(readers, token_type_key):
187184
token_type_embedding = _load_tensor(readers, token_type_key)
188-
assert token_type_embedding.shape == (type_vocab_size, hidden), (
189-
f"Token type embedding shape {token_type_embedding.shape} != "
190-
f"({type_vocab_size}, {hidden})"
191-
)
185+
if token_type_embedding.shape != (type_vocab_size, hidden):
186+
raise ValueError(f'Token type embedding shape {token_type_embedding.shape} != ({type_vocab_size}, {hidden})')
192187
weights["token_type_embedding"] = token_type_embedding.astype(np.float32)
193188
else:
194189
weights["token_type_embedding"] = np.zeros((type_vocab_size, hidden), dtype=np.float32)

‎python/tensorrt_model_connect/families/bloom/plugin.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -62,8 +62,8 @@ def load_weights(
6262
if not _has_tensor(readers, embed_key):
6363
embed_key = "transformer.word_embeddings.weight"
6464
embedding = _load_tensor(readers, embed_key)
65-
assert embedding.shape == (vocab, hidden), (
66-
f"Embedding shape {embedding.shape} != ({vocab}, {hidden})")
65+
if embedding.shape != (vocab, hidden):
66+
raise ValueError(f'Embedding shape {embedding.shape} != ({vocab}, {hidden})')
6767
weights["embedding"] = embedding.astype(np.float32)
6868

6969
# Embedding LayerNorm

‎python/tensorrt_model_connect/families/codegen/plugin.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -59,8 +59,8 @@ def load_weights(
5959

6060
# Token embedding (wte) — no position embedding (uses RoPE)
6161
embedding = _load_tensor(readers, "transformer.wte.weight")
62-
assert embedding.shape == (vocab, hidden), (
63-
f"Embedding shape {embedding.shape} != ({vocab}, {hidden})")
62+
if embedding.shape != (vocab, hidden):
63+
raise ValueError(f'Embedding shape {embedding.shape} != ({vocab}, {hidden})')
6464
weights["embedding"] = embedding.astype(np.float32)
6565

6666
attention_size = hidden

‎python/tensorrt_model_connect/families/convbert/plugin.py‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -97,20 +97,20 @@ def _pfx(key):
9797

9898
# Word embedding
9999
embedding = _load_tensor(readers, _pfx("embeddings.word_embeddings.weight"))
100-
assert embedding.shape == (vocab, embedding_size), (
101-
f"Embedding shape {embedding.shape} != ({vocab}, {embedding_size})")
100+
if embedding.shape != (vocab, embedding_size):
101+
raise ValueError(f'Embedding shape {embedding.shape} != ({vocab}, {embedding_size})')
102102
weights["embedding"] = embedding.astype(np.float32)
103103

104104
# Position embedding
105105
pos_embed = _load_tensor(readers, _pfx("embeddings.position_embeddings.weight"))
106-
assert pos_embed.shape == (max_pos, embedding_size), (
107-
f"Position embedding shape {pos_embed.shape} != ({max_pos}, {embedding_size})")
106+
if pos_embed.shape != (max_pos, embedding_size):
107+
raise ValueError(f'Position embedding shape {pos_embed.shape} != ({max_pos}, {embedding_size})')
108108
weights["position_embedding"] = pos_embed.astype(np.float32)
109109

110110
# Token type embedding
111111
tt_embed = _load_tensor(readers, _pfx("embeddings.token_type_embeddings.weight"))
112-
assert tt_embed.shape == (type_vocab_size, embedding_size), (
113-
f"Token type embedding shape {tt_embed.shape} != ({type_vocab_size}, {embedding_size})")
112+
if tt_embed.shape != (type_vocab_size, embedding_size):
113+
raise ValueError(f'Token type embedding shape {tt_embed.shape} != ({type_vocab_size}, {embedding_size})')
114114
weights["token_type_embedding"] = tt_embed.astype(np.float32)
115115

116116
# Embedding LayerNorm

‎python/tensorrt_model_connect/families/deberta/model/parallel.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -81,7 +81,8 @@ def load_weights(
8181
weights = WeightDict()
8282

8383
embedding = _load_tensor(readers, "deberta.embeddings.word_embeddings.weight")
84-
assert embedding.shape == (vocab, hidden)
84+
if embedding.shape != (vocab, hidden):
85+
raise ValueError(f"Embedding shape {embedding.shape} != ({vocab}, {hidden})")
8586
weights["embedding"] = embedding.astype(np.float32)
8687

8788
if position_biased_input and _has_tensor(

‎python/tensorrt_model_connect/families/deberta/plugin.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -84,7 +84,8 @@ def load_weights(
8484
weights = WeightDict()
8585

8686
embedding = _load_tensor(readers, "deberta.embeddings.word_embeddings.weight")
87-
assert embedding.shape == (vocab, hidden)
87+
if embedding.shape != (vocab, hidden):
88+
raise ValueError(f"Embedding shape {embedding.shape} != ({vocab}, {hidden})")
8889
weights["embedding"] = embedding.astype(np.float32)
8990

9091
if position_biased_input and _has_tensor(

‎python/tensorrt_model_connect/families/deepseek_ocr/plugin.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -121,8 +121,8 @@ def load_weights(
121121

122122
# Embedding
123123
embedding = _load_tensor(readers, f"{lang_prefix}model.embed_tokens.weight")
124-
assert embedding.shape == (vocab, hidden), (
125-
f"Embedding shape {embedding.shape} != ({vocab}, {hidden})")
124+
if embedding.shape != (vocab, hidden):
125+
raise ValueError(f'Embedding shape {embedding.shape} != ({vocab}, {hidden})')
126126
embedding = embedding.astype(io_np_dtype)
127127
weights["embedding"] = embedding
128128

‎python/tensorrt_model_connect/families/deepseek_v2/plugin.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -152,8 +152,8 @@ def load_weights(
152152

153153
# Embedding
154154
embedding = _load_tensor(readers, "model.embed_tokens.weight")
155-
assert embedding.shape == (vocab, hidden), (
156-
f"Embedding shape {embedding.shape} != ({vocab}, {hidden})")
155+
if embedding.shape != (vocab, hidden):
156+
raise ValueError(f'Embedding shape {embedding.shape} != ({vocab}, {hidden})')
157157
weights["embedding"] = embedding.astype(np.float32)
158158

159159
for layer_idx in range(num_layers):

‎python/tensorrt_model_connect/families/distilbert/plugin.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -62,14 +62,14 @@ def load_weights(
6262

6363
# Word embedding
6464
embedding = _load_tensor(readers, "distilbert.embeddings.word_embeddings.weight")
65-
assert embedding.shape == (vocab, hidden), (
66-
f"Embedding shape {embedding.shape} != ({vocab}, {hidden})")
65+
if embedding.shape != (vocab, hidden):
66+
raise ValueError(f'Embedding shape {embedding.shape} != ({vocab}, {hidden})')
6767
weights["embedding"] = embedding.astype(np.float32)
6868

6969
# Position embedding (learned absolute)
7070
pos_embed = _load_tensor(readers, "distilbert.embeddings.position_embeddings.weight")
71-
assert pos_embed.shape == (max_pos, hidden), (
72-
f"Position embedding shape {pos_embed.shape} != ({max_pos}, {hidden})")
71+
if pos_embed.shape != (max_pos, hidden):
72+
raise ValueError(f'Position embedding shape {pos_embed.shape} != ({max_pos}, {hidden})')
7373
weights["position_embedding"] = pos_embed.astype(np.float32)
7474

7575
# DistilBERT has no token_type_embeddings. The encoder builder expects

0 commit comments

Comments
 (0)