load_text_encoder=False training path is unreachable: KeyError 'text' in trainer, null_text_embeds never assigned

[rahulkhunte]

Verified against: a012faa1dc4d71301a7a153c7f9554c081947ea2 (2024-12-22, "Merge pull request #225 from VyneNave/main"), current main at time of writing.

Found while finetuning miniFLUX 384p from precomputed text features.

Summary

train/train_pyramid_flow.py advertises training from precomputed text features: omitting --load_text_encoder sets load_text_fea=True (train/train_pyramid_flow.py:418), the dataset loads .pt feature files instead of raw captions, and PyramidDiTForVideoGeneration.get_text_embeddings takes its else: branch.

That path cannot currently run. Neither defect has surfaced upstream because both shipped training scripts pass --load_text_encoder, so this branch never runs:

$ grep -rn "load_text_encoder" scripts/
scripts/train_pyramid_flow.sh:34:    --load_text_encoder \
scripts/train_pyramid_flow_without_ar.sh:29:    --load_text_encoder \

Bug A stops the run immediately; Bug B is what you hit after working around A.

---

Bug A — the trainer and the dataset disagree about where text features live

trainer_misc/fsdp_trainer.py:85-87 (inside train_one_epoch_with_fsdp, the trainer train/train_pyramid_flow.py:564 actually calls):

                samples = next(data_loader)
                video =  samples['video'].to(accelerator.device)
                text = samples['text']
                identifier = samples['identifier']

But on the precomputed path the dataset emits no 'text' key at all. dataset/dataset_cls.py:188-197 returns the three feature tensors as top-level keys:

            if self.load_text_fea:
                text_fea_path = video_anno['text_fea']
                text_fea = torch.load(text_fea_path, map_location='cpu')
                return {
                    'video': video_latent,
                    'prompt_embed': text_fea['prompt_embed'],
                    'prompt_attention_mask': text_fea['prompt_attention_mask'],
                    'pooled_prompt_embed': text_fea['pooled_prompt_embed'],
                    "identifier": 'video',
                }

samples['text'] therefore raises KeyError: 'text' on the first batch.

There is a three-way disagreement about this contract:

Component	Expects
dataset/dataset_cls.py:188-197	top-level keys, singular: prompt_embed, prompt_attention_mask, pooled_prompt_embed
trainer_misc/fsdp_trainer.py:86	a single samples['text']
pyramid_dit/pyramid_dit_for_video_gen_pipeline.py:606-614	text is a dict with plural keys: prompt_embeds, prompt_attention_mask, pooled_prompt_embeds

Note the singular/plural mismatch between the dataset and the pipeline: even once the trainer is taught to build a text dict, prompt_embed → prompt_embeds and pooled_prompt_embed → pooled_prompt_embeds still need renaming.

Reproduction

Constructs the exact dict the dataset returns, collates it, and runs the two trainer lines. No weights, no dataset, no checkpoint:

import torch
from torch.utils.data.dataloader import default_collate

# Exactly dataset/dataset_cls.py:188-197 when load_text_fea=True — note: no 'text' key.
sample = {
    'video': torch.randn(16, 8, 24, 24),
    'prompt_embed': torch.randn(128, 4096),
    'prompt_attention_mask': torch.ones(128),
    'pooled_prompt_embed': torch.randn(768),
    'identifier': 'video',
}
samples = default_collate([sample] * 4)
print("collated keys:", sorted(samples.keys()))

video = samples['video'].to('cpu')   # fsdp_trainer.py:85
text = samples['text']               # fsdp_trainer.py:86  -> KeyError

Output — reproduced against a clean clone of a012faa just now, not captured from a training run:

collated keys: ['identifier', 'pooled_prompt_embed', 'prompt_attention_mask', 'prompt_embed', 'video']
Traceback (most recent call last):
  File "repro_keyerror_text.py", line 16, in <module>
    text = samples['text']               # fsdp_trainer.py:86  -> KeyError
KeyError: 'text'

Related: device placement

Whatever shape the fix takes, note that dataset_cls.py:190 loads the features with map_location='cpu' and fsdp_trainer.py:85 only transfers video. On the default path this is invisible because self.text_encoder(text, device) (:604) returns device tensors regardless; on the precomputed path the text tensors need an explicit .to(accelerator.device) or they reach the DiT on CPU.

(identifier is a plain str — 'video'/'image', dataset_cls.py:119,196,203,311,372 — collated to a list of strings. It correctly should not be moved to device.)

---

Bug B — self.null_text_embeds is read but never assigned

Once a text dict with the right keys reaches the pipeline, pyramid_dit/pyramid_dit_for_video_gen_pipeline.py:610-612 — the else: branch of get_text_embeddings (def at :598) — reads self.null_text_embeds to implement the 10% CFG dropout:

        else:
            batch_size = len(text['prompt_embeds'])

            for idx in range(batch_size):
                if rand_idx[idx].item():
                    text['prompt_embeds'][idx] = self.null_text_embeds['prompt_embed'].to(device)                   # 610
                    text['prompt_attention_mask'][idx] = self.null_text_embeds['prompt_attention_mask'].to(device)  # 611
                    text['pooled_prompt_embeds'][idx] = self.null_text_embeds['pooled_prompt_embed'].to(device)     # 612

Nothing in the repository ever assigns that attribute. The only three occurrences of the name are the three reads above:

$ grep -rn "null_text_embeds" .
pyramid_dit/pyramid_dit_for_video_gen_pipeline.py:610:                    text['prompt_embeds'][idx] = self.null_text_embeds['prompt_embed'].to(device)
pyramid_dit/pyramid_dit_for_video_gen_pipeline.py:611:                    text['prompt_attention_mask'][idx] = self.null_text_embeds['prompt_attention_mask'].to(device)
pyramid_dit/pyramid_dit_for_video_gen_pipeline.py:612:                    text['pooled_prompt_embeds'][idx] = self.null_text_embeds['pooled_prompt_embed'].to(device)

$ grep -rnE "null_text_embeds\s*=" .
(no matches)

__init__ (:119) sets self.dit, self.text_encoder, self.load_text_encoder (:152), self.vae, … but never self.null_text_embeds.

The crash is probabilistic rather than immediate. rand_idx is built per-row in __call__ at :655:

            rand_idx = torch.rand((batch_size,)) <= self.cfg_rate

and self.cfg_rate is hardcoded to 0.1 at :193 — it is not a constructor argument and is not settable from any script, config file, or CLI flag. So at batch size 4 roughly 34% of steps (1 - 0.9**4) contain at least one dropped row, and a run survives several steps before dying.

Reproduction

Calls the method directly, since Bug A prevents the trainer from ever reaching it. Random tensors only — no weights, no dataset, no checkpoint. Run from the repo root:

import torch
from pyramid_dit.pyramid_dit_for_video_gen_pipeline import PyramidDiTForVideoGeneration as P

runner = object.__new__(P)  # skip __init__: no weights, no VAE, no text encoder
runner.load_text_encoder = False

text = {
    "prompt_embeds":         torch.randn(4, 128, 4096),
    "prompt_attention_mask": torch.ones(4, 128),
    "pooled_prompt_embeds":  torch.randn(4, 768),
}
rand_idx = torch.tensor([0, 1, 0, 0])  # one row picked for CFG dropout

P.get_text_embeddings(runner, text, rand_idx, "cpu")

Traceback — reproduced against a clean clone of a012faa just now, not captured from a training run:

Traceback (most recent call last):
  File "repro_null_text_embeds.py", line 14, in <module>
    P.get_text_embeddings(runner, text, rand_idx, "cpu")
  File ".../torch/utils/_contextlib.py", line 124, in decorate_context
    return func(*args, **kwargs)
  File ".../Pyramid-Flow/pyramid_dit/pyramid_dit_for_video_gen_pipeline.py", line 610, in get_text_embeddings
    text['prompt_embeds'][idx] = self.null_text_embeds['prompt_embed'].to(device)
AttributeError: 'PyramidDiTForVideoGeneration' object has no attribute 'null_text_embeds'

Note on the expected shapes

Lines 610-612 assign into a single batch row ([idx]), so whatever null_text_embeds holds must be unbatched — for miniFLUX, prompt_embed [128, 4096], prompt_attention_mask [128], pooled_prompt_embed [768]. Supplying batched [1, 128, 4096] tensors trades the AttributeError for a shape mismatch, so this is easy to get subtly wrong.

---

Why I haven't sent a patch

Both fixes involve a design decision I'd rather not make on your behalf.

For Bug A, which of the three components is authoritative? The trainer could assemble a text dict from the top-level keys (smallest change, one file); or the dataset could nest them under 'text' with the plural names the pipeline expects (keeps the trainer generic, changes the dataset's output contract). Either way the singular/plural rename has to land somewhere.

For Bug B, where should the null embedding come from? The options I see:

1. Caller passes it in (constructor arg or setter) — whoever precomputed the dataset supplies a null embedding produced by the same encoder, which is the only way to guarantee they match. Adds a required argument to the precomputed path.
2. Load it from the feature directory as a sibling artifact, so it travels with the dataset it must match. No new argument, but introduces a file-layout convention.
3. Build it in __init__ by loading the text encoder just long enough to encode "" — self-contained, but loads the encoder in exactly the configuration whose purpose is to avoid loading it.
4. Zeros of the right shape — cheapest and needs no encoder, but is not equivalent to the encoding of "" and would make CFG behave differently from the load_text_encoder=True path.

Happy to send PRs for whichever shape you prefer.