Skip to content

Remove DictWrapper class from LLM model#4096

Open
w4nderlust wants to merge 2 commits intomainfrom
remove-dictwrapper
Open

Remove DictWrapper class from LLM model#4096
w4nderlust wants to merge 2 commits intomainfrom
remove-dictwrapper

Conversation

@w4nderlust
Copy link
Copy Markdown
Collaborator

Summary

  • Remove the DictWrapper class from ludwig/models/llm.py and replace it with a plain Python dict
  • DictWrapper existed to prevent input/output features from being registered as nn.Module submodules (which would confuse DeepSpeed). A plain dict achieves the same goal more simply since it is not an nn.Module
  • Fix items()[0] indexing (which relied on LudwigFeatureDict.items() returning a list) to use next(iter(...)) compatible with dict_items views

Test plan

  • Verify LLM model initialization works correctly
  • Verify LLM fine-tuning works (features should not appear as DeepSpeed-managed modules)
  • Verify LLM prediction/generation works
  • Run existing LLM unit tests

Replace the DictWrapper class with a plain Python dict in
LLM.create_feature_dict(). The DictWrapper existed solely to prevent
input/output features from being registered as nn.Module submodules
(which would confuse DeepSpeed). A plain dict achieves the same goal
more simply -- it is not an nn.Module, so PyTorch and DeepSpeed never
see it.

Key changes:
- Remove the DictWrapper class (42 lines)
- Remove unused LudwigFeatureDict import
- create_feature_dict() now returns {}
- Fix items()[0] indexing (list->dict_items view) with next(iter(...))
@github-actions
Copy link
Copy Markdown

github-actions bot commented Apr 7, 2026

Test Results

   10 files  ±0     10 suites  ±0   1h 46m 4s ⏱️ - 1m 23s
3 619 tests ±0  3 589 ✅ ±0  30 💤 ±0  0 ❌ ±0 
3 708 runs  ±0  3 665 ✅ ±0  43 💤 ±0  0 ❌ ±0 

Results for commit e2e28eb. ± Comparison against base commit 7aa7353.

♻️ This comment has been updated with latest results.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant