HuggingFace Integration Friction: Weight Tying, Generation, Trainer, and Dataset Collator Alignment

[david-thrower]

Background

The avant-garde model core (recurrent heterogeneous graph, LTI, ACT, Mamba-2 SSD, Titans) is stable and working. The HF wrapper layer (hf_model.py, trainer.py) has accumulated friction points that break standard HF workflows (Trainer, pipeline, AutoModelForCausalLM.from_pretrained, model.generate()). A previous patch fixed the redundant lm_head dead-weight bug; this issue completes the V5 integration.

Problems Identified

#	Problem	Location	Impact
1	Weight tying bypasses HF standard mechanism	hf_model.py HelixPreTrainedModel	_tied_weights_keys = {} and get_expanded_tied_weights_keys override kill model.tie_weights() and from_pretrained weight restoration.
2	Custom generate_ext() instead of standard generate()	hf_model.py HelixForCausalLM	Cannot use pipeline("text-generation", ...) or standard StoppingCriteria; agent must re-implement stop strings, temperature, top-p manually.
3	Custom Trainer reinvents transformers.Trainer	trainer.py	Duplicates AMP, gradient accumulation, scheduler, logging. Loses DeepSpeed/FSDP support.
4	Dataset overlap masking at risk from standard collators	dataset.py / trainer.py	DataCollatorForLanguageModeling regenerates labels and destroys custom -100 overlap masks.
5	Tokenizer wrapper adds friction	tokenizer.py	HelixTokenizer is not a first-class PreTrainedTokenizer; passing it to HF Trainer requires ._backend indirection.
6	Auto-registration wrapped in silent try/except	hf_model.py bottom	Registration failures are swallowed.

Architectural Constraints (Non-Negotiable)

The following are sacred and must not be modified:

• helix_lm/graph.py — heterogeneous graph wiring
• helix_lm/recurrent.py — LTI injection, ACT halting, recurrence
• helix_lm/nodes.py — all node implementations (attention variants, SwiGLU, SSM, Titans, gates)
• helix_lm/mamba2.py — Mamba-2 SSD parallel scan
• helix_lm/rope.py — rotary embeddings

The following behaviors must be preserved exactly:

• Document-aware chunking with overlap masking: DocumentAwareDataset must continue to return labels with -100 on overlap heads and padding tails. Standard DataCollatorForLanguageModeling must not be used; it overwrites these labels.
• No KV-cache / no past_key_values state: The recurrent graph re-initializes node_states = {} on every forward call. Generation must pass the full sequence (or last seq_len window) on every step, never just input_ids[:, -1:]. config.use_cache must remain False.
• Parameter count stability: After all fixes, HelixForCausalLM(HelixConfig.tiny(vocab_size=50257)) must report exactly 13,347,974 parameters.

Proposed Fix Stages

See attached agent prompt package for executable stage-by-stage tasks.

[david-thrower]

Reserved for work on this (Agent instructions, needs reviewed before submitted for completion)

Stage 0: Code Examination & Constraint Mapping

Purpose: Build a mental model of the integration surface without touching anything.

Read these files (do not modify):

• helix_lm/hf_model.py
• helix_lm/model.py
• helix_lm/trainer.py
• helix_lm/dataset.py
• helix_lm/tokenizer.py
• helix_lm/config.py
• helix_lm/__init__.py
• quick_demo_cpu.py
• examples/train_hf_full.py
• requirements.txt

Answer these questions in your response before proceeding:

1. Weight tying: Where is _tied_weights_keys defined, and what value does it have? What does get_expanded_tied_weights_keys return?
2. Generation: What is the exact signature of generate_ext? List every argument. What is the exact logic for stop-string detection?
3. Trainer: What does the custom Trainer class do that transformers.Trainer does not? List at least 3 unique behaviors.
4. Dataset: What keys does DocumentAwareDataset.__getitem__ return? What values are -100 and why?
5. Core architecture: What is the exact path from HelixForCausalLM.__init__ to the active output projection during forward()? (Trace the tensor through self.model.* layers.)
6. Registration: How are AutoConfig and AutoModelForCausalLM registered? Is the registration silent-failing?

If you cannot answer any question with certainty, re-read the file. Do not proceed until all six are answered.

---

Stage 1: Weight Tying & Model Registration

Goal: Make save_pretrained / from_pretrained / pipeline restore weights correctly.

Modify:

• helix_lm/hf_model.py

Do NOT modify:

• Any file in helix_lm/graph.py, recurrent.py, nodes.py, mamba2.py, rope.py

Specific changes:

1. In HelixPreTrainedModel:

   • DELETE the line _tied_weights_keys = {}
   • DELETE the method get_expanded_tied_weights_keys
   • Keep _init_weights, to_device, _resolve_device unchanged.

2. In HelixForCausalLM:

   • Add class attribute: _tied_weights_keys = ["lm_head.weight"]
   • Ensure get_input_embeddings(self) returns self.model.embed
   • Ensure set_input_embeddings(self, value) sets self.model.embed = value
   • Ensure get_output_embeddings(self) returns self.lm_head
   • Ensure set_output_embeddings(self, new_embeddings) sets self.lm_head = new_embeddings
   • Verify post_init() is called in __init__ (it already should be). After this change, post_init() will automatically tie lm_head.weight to embed.weight when config.tie_word_embeddings=True.

3. At bottom of hf_model.py:

   • Replace the try/except block with direct registration (no silent swallowing):

   AutoConfig.register("helix", HelixConfig)
   AutoModelForCausalLM.register(HelixConfig, HelixForCausalLM)

Test command:

python -c "
from helix_lm import HelixConfig, HelixForCausalLM
cfg = HelixConfig.tiny(vocab_size=50257, tie_word_embeddings=True)
m = HelixForCausalLM(cfg)
print('Tied?', m.lm_head.weight is m.get_input_embeddings().weight)
print('Count:', m.count_parameters()['total'])
assert m.count_parameters()['total'] == 13347974, 'Param count drift!'
"

Expected output: Tied? True and Count: 13347974 (or 13347974 for untied; verify based on config default).

---

Stage 2: Generation Migration (generate_ext → standard generate)

Goal: Replace the custom generate_ext method with standard generate() plus a StoppingCriteria class. Preserve all existing generation behavior (temperature, top-k, top-p, repetition penalty, stop strings).

Modify:

• helix_lm/hf_model.py

Do NOT modify:

• Core architecture files
• quick_demo_cpu.py (yet; you'll update it in Stage 5)

Specific changes:

1. Create StopStringCriteria inside hf_model.py (or a new helix_lm/generation_utils.py if you prefer, but keep it importable from hf_model.py).

   Requirements:

   • Inherit from transformers.StoppingCriteria
   • __init__(self, tokenizer, stop_strings: List[str])
   • __call__(self, input_ids: torch.Tensor, scores: torch.Tensor, **kwargs) -> bool
   • Decode each batch element with skip_special_tokens=True
   • Return True if any stop string appears in any batch element

2. DELETE the generate_ext method from HelixForCausalLM.

3. UPDATE prepare_inputs_for_generation:

   • The current code trims to input_ids[:, -1:] when past_key_values is not None.
   • Because our model has no real KV-cache, this trim is dangerous.
   • Change it to: Always return the full input_ids (or truncated to config.seq_len if longer). Never trim to last token.

   def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs):
       # Never trim to last token; our model recomputes from full sequence
       if input_ids.size(1) > self.config.seq_len:
           input_ids = input_ids[:, -self.config.seq_len:]
       return {
           "input_ids": input_ids,
           "attention_mask": kwargs.get("attention_mask"),
           "use_cache": False,
       }

4. Ensure config.use_cache = False in HelixConfig defaults (or set it explicitly during generation).

5. Preserve _reorder_cache as a no-op returning past_key_values unchanged.

Test command:

python -c "
from helix_lm import HelixConfig, HelixForCausalLM, HelixTokenizer
from transformers import StoppingCriteriaList

cfg = HelixConfig.tiny(vocab_size=50257, use_cache=False)
model = HelixForCausalLM(cfg).eval()
tok = HelixTokenizer('gpt2')

prompt = 'In 1492,'
ids = tok.encode(prompt, return_tensors='pt')
out = model.generate(
    ids,
    max_new_tokens=25,
    temperature=0.8,
    top_k=50,
    top_p=0.95,
    stopping_criteria=StoppingCriteriaList([
        model.__class__.__module__  # you will define StopStringCriteria properly
    ])  # actually test without stopping criteria first, then with
)
print(tok.decode(out[0]))
"

Also run existing tests:

python smoke_test_short.py
python smoke_test_long.py

---

Stage 3: Trainer Subclassing (Custom → transformers.Trainer)

Goal: Migrate the custom Trainer in trainer.py to subclass transformers.Trainer. Preserve NaN/Inf skipping, generation samples, and lazy dataset support. Eliminate redundant AMP/scheduler/accumulation logic.

Modify:

• helix_lm/trainer.py

Do NOT modify:

• helix_lm/dataset.py (except to add __len__ compatibility if missing)
• Core model files

Specific changes:

1. Import from transformers import Trainer as HFTrainer

2. Redefine class Trainer(HFTrainer):

3. Override compute_loss to add NaN/Inf detection:

   • Call super().compute_loss(...) to get the base loss
   • If loss is NaN or Inf, return 0.0 * loss (creates a dummy backward pass) and log a warning
   • Do NOT skip the step by returning None; HF Trainer expects a tensor. Instead, return a zero-loss tensor so gradients are zero but step count advances.

4. Override evaluate to preserve live generation samples:

   • Call super().evaluate(...) to get base metrics
   • If self.args.should_save or epoch-end, run 3 example prompts through model.generate(...) (standard generate, not generate_ext)
   • Print samples to stdout (or logger)

5. Override get_train_dataloader and get_eval_dataloader ONLY if necessary to preserve DocumentAwareDataset lazy initialization. If standard HF Trainer can already consume your dataset (it returns dicts of tensors), you may not need to override these.

6. Remove redundant methods from your old custom Trainer:

   • train_epoch (HF Trainer handles this)
   • Manual AMP / GradScaler logic (HF Trainer handles via TrainingArguments.fp16/bf16)
   • Manual scheduler stepping (HF Trainer handles via optim arg)
   • Manual gradient accumulation (HF Trainer handles via gradient_accumulation_steps arg)
   • Manual logging / progress bar (HF Trainer handles via logging_steps, report_to)

7. Update quick_demo_cpu.py to use TrainingArguments + your new Trainer:

   from transformers import TrainingArguments
   from helix_lm import Trainer  # now subclasses HF Trainer
   
   training_args = TrainingArguments(
       output_dir="./checkpoints",
       num_train_epochs=25,
       per_device_train_batch_size=8,
       learning_rate=3e-4,
       weight_decay=0.1,
       warmup_steps=100,
       max_grad_norm=1.0,
       logging_steps=10,
       eval_strategy="epoch",
       save_strategy="epoch",
       fp16=False,
   )
   
   trainer = Trainer(
       model=model,
       args=training_args,
       train_dataset=train_dataset,  # DocumentAwareDataset
       eval_dataset=val_dataset,
       data_collator=data_collator,  # DataCollatorWithPadding (Stage 4)
       tokenizer=tokenizer._backend,
   )
   trainer.train()

Test command:

python quick_demo_cpu.py

Must complete 2 epochs, show stable parameter count, and generate text without errors.

---

Stage 4: Dataset & Collator Alignment

Goal: Ensure DocumentAwareDataset works with HF Trainer and the correct data collator. Prevent label destruction.

Modify:

• helix_lm/trainer.py (collator reference in constructor)
• examples/train_hf_full.py and quick_demo_cpu.py (collator usage)

Do NOT modify:

• helix_lm/dataset.py logic (preserving overlap masking)

Specific changes:

1. Use DataCollatorWithPadding:

   from transformers import DataCollatorWithPadding
   data_collator = DataCollatorWithPadding(tokenizer=tokenizer._backend)

   Pass this to Trainer.

2. Verify DocumentAwareDataset is compatible with HF Trainer:

   • It must implement __len__
   • __getitem__ must return a dict where all values are either tensors or simple types that can be batched
   • The is_natural_stop bool tensor might cause issues with DataCollatorWithPadding because it tries to pad everything. If so, do not modify the dataset. Instead, write a minimal custom collate function:

   def collate_fn(examples):
       # Separate tensor fields from meta fields
       tensor_keys = ["input_ids", "labels", "attention_mask"]
       meta_keys = ["is_natural_stop"]  # bool, no padding needed
      
       batch = tokenizer._backend.pad(
           [{k: v for k, v in ex.items() if k in tensor_keys} for ex in examples],
           padding=True,
           return_tensors="pt",
       )
       batch["is_natural_stop"] = torch.stack([ex["is_natural_stop"] for ex in examples])
       return batch

   Use this collate function instead of DataCollatorWithPadding if the standard one chokes on bools.

3. Explicitly forbid DataCollatorForLanguageModeling in comments and code:

   # NOTE: Do NOT use DataCollatorForLanguageModeling here.
   # It regenerates labels from input_ids and destroys our custom overlap masking.

Test command:

python -c "
from helix_lm.dataset import DocumentAwareDataset
from helix_lm.tokenizer import HelixTokenizer

tok = HelixTokenizer('gpt2')
ds = DocumentAwareDataset(
    texts=['Once upon a time there was a cat.'] * 5,
    tokenizer=tok,
    seq_len=32,
    stride=16,
    min_tail_len=1,
)

sample = ds[0]
print('Keys:', sample.keys())
print('Labels contain -100:', (sample['labels'] == -100).any().item())
print('is_natural_stop:', sample['is_natural_stop'])
assert 'input_ids' in sample
assert 'labels' in sample
assert 'attention_mask' in sample
"

---

Stage 5: Tokenizer Simplification (Optional but Recommended)

Goal: Reduce friction when passing tokenizers to HF Trainer / pipeline.

Modify:

• helix_lm/tokenizer.py (add HF-native methods)
• quick_demo_cpu.py (use tokenizer._backend where HF objects are expected)

Do NOT modify:

• The char-level vocab builder logic (keep it working)

Specific changes:

1. In HelixTokenizer, add passthrough methods so it quacks like an HF tokenizer:

   @property
   def pad_token_id(self):
       return self._backend.pad_token_id if self._backend else 0
   
   def pad(self, *args, **kwargs):
       return self._backend.pad(*args, **kwargs)

   This lets DataCollatorWithPadding(tokenizer=helix_tokenizer) work directly without ._backend.

2. In training scripts, prefer passing the HelixTokenizer instance directly (if passthrough works) or keep ._backend with a TODO comment.

Test command:

python -c "
from helix_lm.tokenizer import HelixTokenizer
tok = HelixTokenizer('gpt2')
print('pad_token_id:', tok.pad_token_id)
print('vocab_size:', len(tok))
"

---

Stage 6: Integration Smoke Test

Goal: Verify the entire stack works end-to-end with standard HF APIs.

Create: smoke_test_hf_integration.py

Requirements:

1. Instantiate HelixForCausalLM with HelixConfig.tiny(vocab_size=50257, use_cache=False, tie_word_embeddings=True)
2. Confirm parameter count is 13,347,974
3. Instantiate HelixForCausalLM with HelixConfig.tiny(vocab_size=50257, use_cache=False, tie_word_embeddings=False)
4. Confirm parameter count is 13,347,974 (untied; verify this is expected)
5. Run model.generate() with StoppingCriteriaList([StopStringCriteria(...)])
6. Verify stop string works (generation stops before max_new_tokens if stop string encountered)
7. Save to temp dir with model.save_pretrained()
8. Reload with AutoModelForCausalLM.from_pretrained(temp_dir)
9. Verify reloaded model produces identical logits on same input
10. If tie_word_embeddings=True, verify lm_head.weight is model.get_input_embeddings().weight after reload