google-research / google-research/tabfm

config.json keys are merged into the model constructor unvalidated (OOM kill, ZeroDivisionError)

Open
#91 0 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.