_from_pretrained copies every key from a model directory's config.json into model_kwargs, with no allowlist and no validation (tabfm/src/pytorch/tabfm_v1_0_0.py:60-70):
def _apply_config(cfg):
if "is_classifier" not in model_kwargs and "task" in cfg:
model_kwargs["is_classifier"] = cfg.pop("task") == "classification"
for key in ("model_type", "version", "framework"):
cfg.pop(key, None)
for k, v in cfg.items():
if k not in model_kwargs:
model_kwargs[k] = vThose keys land directly on the constructor (tabfm/src/pytorch/model.py:695-714), where several of them size allocations:
ff = embed_dim * ff_factor
icl_dim = embed_dim * row_num_cls
self.cls_tokens = nn.Parameter(torch.zeros(row_num_cls, embed_dim))
No constructor in the tree validates its dimensional arguments — I searched for raise ValueError/assert against embed_dim, nhead, max_classes, ff_factor, row_num_cls and found none on the PyTorch side.
Reproduced on b15593e4c1111ddb5f4f30dd2957df2edbaa04ca, container capped at 3 GB, loading a local directory via TabFM_HF.from_pretrained(dir):
config.json = {"task": "classification", "embed_dim": 1048576, "ff_factor": 1024, "row_num_cls": 16384}
-> Killed, exit 137 (SIGKILL, out of memory)
max_classes does the same on its own, and a zero head count gives a less obvious failure:
[torch max_classes=2**28] Killed, exit 137
[torch nhead=0] ZeroDivisionError: integer division or modulo by zero
[torch embed_dim=-1] RuntimeError: Trying to create tensor with negative dimension <- handled well
I want to be careful not to overstate this. Loading a model directory is a trust decision, the serious case (code execution through pickled weights) is already closed by weights_only=True in huggingface_hub, and the worst outcome here is that the loading process dies. So I would call it robustness rather than a vulnerability, and I am not asking for a security response.
It seems worth a modest fix regardless: validate the dimensional keys where they arrive, or accept only a known set of config keys instead of merging everything. It would also turn nhead: 0 into a message that says what is wrong. This matters a little more given #88 — the weights are fetched from an unpinned ref, so config.json is not fully under your control at load time.
Disclosure: I used an AI assistant to help find this. I ran the reproductions myself.
_from_pretrainedcopies every key from a model directory'sconfig.jsonintomodel_kwargs, with no allowlist and no validation (tabfm/src/pytorch/tabfm_v1_0_0.py:60-70):Those keys land directly on the constructor (
tabfm/src/pytorch/model.py:695-714), where several of them size allocations:No constructor in the tree validates its dimensional arguments — I searched for
raise ValueError/assertagainstembed_dim,nhead,max_classes,ff_factor,row_num_clsand found none on the PyTorch side.Reproduced on
b15593e4c1111ddb5f4f30dd2957df2edbaa04ca, container capped at 3 GB, loading a local directory viaTabFM_HF.from_pretrained(dir):max_classesdoes the same on its own, and a zero head count gives a less obvious failure:I want to be careful not to overstate this. Loading a model directory is a trust decision, the serious case (code execution through pickled weights) is already closed by
weights_only=Trueinhuggingface_hub, and the worst outcome here is that the loading process dies. So I would call it robustness rather than a vulnerability, and I am not asking for a security response.It seems worth a modest fix regardless: validate the dimensional keys where they arrive, or accept only a known set of config keys instead of merging everything. It would also turn
nhead: 0into a message that says what is wrong. This matters a little more given #88 — the weights are fetched from an unpinned ref, soconfig.jsonis not fully under your control at load time.Disclosure: I used an AI assistant to help find this. I ran the reproductions myself.