[Bug]: doctor reports "no tensor fills these core roles: token embedding, output head" for a healthy DeepSeek V4 container

[pauloapoloni]

Commit or version

v1.11.0 (HEAD a8f2ca6)

Environment

Ubuntu 24.04, gcc 13.3, 16 threads, 30 GB RAM, NVMe ext4. CPU-only build of the V4 engine (make deepseek-v4).
Model: puwaer/DeepSeek-V4-Flash-0731-reap-150b (REAP-pruned DeepSeek V4 Flash, 84.7 GB, 43 layers x 132 experts, official fp4 experts / fp8 dense, no conversion).
Python used by coli: the same interpreter that runs the launcher.

Reproduction steps

hf download puwaer/DeepSeek-V4-Flash-0731-reap-150b --local-dir /mnt/Data/models/DeepSeek-V4-Flash-reap-150b
cd /home/apoloni/repos/colibri
python3 c/coli doctor --model /mnt/Data/models/DeepSeek-V4-Flash-reap-150b --deep

Expected behavior

model.required passes: the checkpoint contains a token embedding and an output head, and the engine loads and generates with them.

Actual behavior and logs

$ python3 c/coli doctor --model /mnt/Data/models/DeepSeek-V4-Flash-reap-150b --deep
[fail] model.required     no tensor fills these core roles: token embedding, output head
[  ok] model.index        model index matches every scanned tensor
[skip] storage.mirror     no mirror directory is configured
...
result error

The checkpoint is not missing anything. Its index has 35,620 tensors, including exactly the two names the engine looks up:

$ python3 -c "import json;print([n for n in json.load(open('model.safetensors.index.json'))['weight_map'] if n.endswith('.weight') and ('embed' in n or 'head' in n)][:6])"
['embed.weight', 'head.weight', 'hc_head_base', 'hc_head_fn', 'hc_head_scale']

and c/deepseek_v4.c finds them by those names:

c/deepseek_v4.c:11515:  const ColiSafetensorsTensor *embed = coli_st_find(index, "embed.weight");
c/deepseek_v4.c:1196:   *head = coli_st_find(index, "head.weight");
c/deepseek_v4.c:11629:  const ColiSafetensorsTensor *head = coli_st_find(index, "head.weight");

The engine runs this container end to end (CPU path verified: prompt "Say hello in exactly one short sentence." produces Hello.<end_of_sentence>), so the doctor verdict disagrees with the engine that actually loads the file.

Root cause (in c/doctor.py)

The role predicates match on names taken from other families:

def _is_embedding(name):
    return name.endswith("embed_tokens.weight")

def _is_output_head(name):
    return name.endswith("lm_head.weight")

DeepSeek V4 names them embed.weight and head.weight, so both roles are reported missing while model.index — which scans the same tensors — is green. _is_final_norm happens to survive here only because the container has a matching norm.weight.

This is the same class of defect as the #1365 case documented in the comment directly above those predicates ("the doctor called a healthy model broken"). The comment says the fix was to match on the tail rather than a full prefix; these two still match on a full literal suffix that no DeepSeek V4 container uses.

Suggested fix

Either extend the two predicates with the DeepSeek V4 spellings (endswith("embed.weight") / endswith("head.weight"), with the existing exclusions kept so hc_head_* cannot stand in for an output head), or derive the required roles from what each engine actually looks up — deepseek_v4.c is the authority for this family and already states both names.

Happy to test a patch: I have this container locally and can run doctor --deep and a generation before and after.

[JustVugg]

Confirmed, and your root cause is exactly right. I checked deepseek_v4.c and it looks those two names up verbatim at four call sites (embed.weight at 11624 and 18227, head.weight at 1196, 11761 and 11842), so the engine and the doctor genuinely disagree about the same file.

There is an irony worth recording in the fix. The comment sitting directly above those two predicates documents #1365, which is this defect. The lesson taken from it was to match the tail rather than the whole name, and that made the predicates prefix-agnostic. It did not make them name-agnostic, and DeepSeek V4 does not spell the roles the way the families in front of it did. _is_final_norm survived only because .norm.weight happens to be a spelling V4 shares.

On the suggested fix

Option one is the right size, but the plain endswith form has a false positive it is worth not shipping. "pos_embed.weight".endswith("embed.weight") is True, and so is the vision tower's patch_embed.weight:

with the endswith form proposed above
  pos_embed.weight              embedding=True
  vision.patch_embed.weight     embedding=True

That is the same shape as the case _is_final_norm already guards against, where model.visual.post_layernorm.weight would otherwise stand in for a final norm that is not there. A positional embedding standing in for a token embedding would turn this bug into its mirror image: a checkpoint that really is missing its embedding would pass.

Comparing the component rather than the tail fixes both directions. The part immediately before .weight has to BE the role:

def _core_component(name, spellings):
    parts = name.split(".")
    return (len(parts) >= 2 and parts[-1] == "weight"
            and parts[-2] in spellings and ".layers." not in name)

def _is_embedding(name):   return _core_component(name, {"embed_tokens", "embed"})
def _is_output_head(name): return _core_component(name, {"lm_head", "head"})

Checked against the names in play:

embed.weight                              embedding=True
head.weight                                              head=True
model.embed_tokens.weight                 embedding=True
lm_head.weight                                           head=True
model.language_model.embed_tokens.weight  embedding=True
pos_embed.weight                          (neither)
vision.patch_embed.weight                 (neither)
hc_head_base / hc_head_fn / hc_head_scale (neither)

hc_head_* falls out for free because those do not end in .weight, so your concern about them standing in for the output head is covered without an explicit exclusion. The .layers. exclusion is borrowed from _is_final_norm for the same reason it exists there: a head or an embedding inside the layer stack is not the model's.

On your option two

Deriving the roles from what each engine looks up is the right instinct, and it is what would stop this recurring a fourth time. It is not reachable today: the engines resolve these names in C at load time, several of them with fallbacks (qwen38.c probes the nested spelling first and the flat one second), so there is no table the doctor can read. A named set per role is the honest middle: still prefix-agnostic, and it states the spellings that exist rather than assuming one family's.

I will put up a patch with this and a V4 case added to tests/test_doctor.py, which already covers the tied-embeddings and role-naming paths. Your offer to test is the part we cannot do here: nobody on this side has that container, and "doctor is green and the engine still generates" is the check that matters, so I will ping you when the branch is up.

[JustVugg]

Patch is up: #1608, branch fix/doctor-core-roles-v4.

It takes the component form rather than the tails, for the pos_embed.weight reason above, and it also moves _is_final_norm to the same shape. That third one was not what you reported, but the test found it: the old .norm.weight tail requires a prefix, so a bare norm.weight at the root does not match. A container that names its roles flat, which is what V4 does with the other two, would have failed the third role for the same reason it failed the first two. Your container clearly does not, since the report shows that role filled, but it was sitting there.

When you have a moment, the two runs that would settle it on your box:

git fetch origin fix/doctor-core-roles-v4 && git checkout fix/doctor-core-roles-v4
python3 c/coli doctor --model /mnt/Data/models/DeepSeek-V4-Flash-reap-150b --deep

and then a generation on the same container, so we have doctor green and the engine still producing text from the same checkout rather than just the first half.