google-research / google-research/tabfm
config.json keys are merged into the model constructor unvalidated (OOM kill, ZeroDivisionError)
- Dominant language
- Python
- Stars
- 2.6k
- Forks
- 270
- Avg merge
- 1d 7h
- Merged PRs (30d)
- 1
Description
`_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`):
```python
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] = v
```
Those keys land directly on the constructor (`tabfm/src/pytorch/model.py:695-714`), where several of them size allocations:
```python
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.
Contributor guide
Research direction
Start with _apply_config in tabfm/src/pytorch/tabfm_v1_0_0.py:60-70 and the constructor in tabfm/src/pytorch/model.py:695-714; reproduce loading a local directory with the malformed values shown in the issue. The work is done when invalid dimensional configuration is rejected with a clear error before unsafe allocations or division, while valid configuration loading still works.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- backend
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 58/100