Add bd3lm - #54
Add bd3lm#54
Conversation
Removed note about compatibility of newer released repos.
Review: pretrained checkpoint loading vs LLaDA/Dream, plus a generation smoke testLooked specifically at (1) whether loading 1. Loading path: this is a new pathway for the same objectiveLLaDA and Dream load released Hub weights through the shared xLM inference/train loader:
BD3LM does not use that path. It adds a model-local loader:
That prefix strip is a real difference versus LLaDA/Dream. The released checkpoint is an HF After
The in- 2. Tokenizer mismatch on the advertised OWT fine-tune pathThe released checkpoint is GPT-2 + one mask token:
3. Inference smoke test (
|
Added instructions for unconditional generation using bd3lm.
Made the tokenizer match the released checkpoints and added an inference config for unconditional generation. Ran gen. ppl on kuleshov-group/bd3lm-owt-block_size4 at length 1024 and got 25.74 over 300 samples, against 25.70 in the paper.
Added instructions for converting and generating with the released checkpoint.
Clarify usage instructions for converting HF checkpoints.
Updated docstring to clarify loss computation logic.
Removed redundant explanation about token rows in YAML.
Updated various sections for clarity and consistency, including loss handling and config details.
Clarified descriptions of `indices` and `var_length` parameters. Updated instructions for using `eval.model_only_checkpoint_path` and added details about verification conditions.
Moved the transformer under self.backbone so the released checkpoint names match ours 1:1, and hub.repo_id can load them directly like LLaDA does.
Generation is now one command with +pretrained=auto, and the converter is gone.
Removed redundant explanations about diffusion denoising, loss on padding, and predictor behavior. Updated sections for clarity and conciseness.
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
| # @package _global_ | ||
| # Load the released OWT checkpoint matching your block_size. | ||
| # | ||
| # xlm job_type=train ... +pretrained=auto # block_size=4 -> block_size4 | ||
| # xlm job_type=eval ... +pretrained=auto block_size=8 # -> block_size8 | ||
| # | ||
| hub: | ||
| repo_id: kuleshov-group/bd3lm-owt-block_size${block_size} |
There was a problem hiding this comment.
Please rename this overlay from pretrained=auto to pretrained=kuleshov_group_bd3lm (file kuleshov_group_bd3lm.yaml). auto hides which checkpoint family this is; the name should say it loads the kuleshov-group OWT releases.
Also restrict it to OWT (owt_bd3lm / owt_bd3lm_inference). Do not use it on star.
|
|
||
| own = self.state_dict() | ||
| mismatched = [] | ||
| for key in [k for k in state_dict if k.startswith(prefix)]: | ||
| tail = key[len(prefix):] | ||
| if tail in own and own[tail].shape != state_dict[key].shape: | ||
| mismatched.append((tail, tuple(own[tail].shape), tuple(state_dict.pop(key).shape))) | ||
| if mismatched: | ||
| log = logging.getLogger(__name__).warning | ||
| log("[bd3lm] %d tensor(s) skipped on shape and will train from scratch - " | ||
| "usually a vocabulary difference:", len(mismatched)) | ||
| for name, ours, theirs in mismatched: | ||
| log("[bd3lm] %s: model %s vs checkpoint %s", name, ours, theirs) |
There was a problem hiding this comment.
Please drop this shape-mismatch skip. We should not support loading kuleshov-group/bd3lm-owt-* onto star (or any non-GPT-2 vocab). Star trains from scratch; transferring English LM blocks while randomly initializing embed/lm-head is not a path we want.
Keep the sampling_eps_* drop above if a strict OWT hub.repo_id load still needs it. Just do not special-case a 20-vs-50258 vocab.
| ## Fine-tuning from a released checkpoint | ||
|
|
||
| Add one flag: | ||
|
|
||
| ```bash | ||
| xlm job_type=train job_name=my_finetune \ | ||
| experiment=star_medium_bd3lm \ | ||
| +pretrained=auto | ||
| ``` | ||
|
|
||
| The checkpoint is derived from `block_size`, so `block_size=8` pulls the block_size8 | ||
| weights. Four are compatible, all 768 / 12 blocks / 12 heads (`bd3lm_small`) on the | ||
| GPT-2 vocabulary: | ||
|
|
||
| | block size | checkpoint | | ||
| |---|---| | ||
| | 4 | `kuleshov-group/bd3lm-owt-block_size4` | | ||
| | 8 | `kuleshov-group/bd3lm-owt-block_size8` | | ||
| | 16 | `kuleshov-group/bd3lm-owt-block_size16` | | ||
| | 1024 | `kuleshov-group/bd3lm-owt-block_size1024-pretrain` | | ||
|
|
||
| On a task whose vocabulary is not GPT-2's, the three vocabulary-sized tensors cannot | ||
| transfer. They are skipped and reported, and the transformer blocks still load: | ||
|
|
||
| ``` | ||
| [bd3lm] 3 tensor(s) skipped on shape and will train from scratch - usually a vocabulary difference: | ||
| [bd3lm] backbone.vocab_embed.embedding: model (27, 768) vs checkpoint (50258, 768) |
There was a problem hiding this comment.
Please remove this star + +pretrained=auto fine-tune path. Star should not load the released OWT checkpoint.
Document the overlay only on OWT generation / continued LM training, and use +pretrained=kuleshov_group_bd3lm once the config is renamed. Same change in docs/models/bd3lm.md (the matching “Fine-tuning from a released checkpoint” section).
There was a problem hiding this comment.
Fixes:
Removed the star fine-tuning on released OWT checkpoints
Renamed auto.yaml to kuleshov_group_bd3lm.yaml
Removed loading the released checkpoint for star, and renamed the pre-trained yaml used for OWT inference.
Removed NaN checks and associated error handling from transformer block.
Removed NaN check and associated print statements before processing blocks.
Removed NaN loss debug checks from loss calculation.
Removed NaN check from model output.
Removed unnecessary blank lines to improve code readability.
Added entropy-based stopping criteria to model configuration.
Added model configuration for sampling entropy control.
Added model configuration with sampling settings.
Removed commented-out code for clarity.
Refactor entropy computation to handle each row individually and update stopping conditions.
Refactor _compute_entropy method to compute entropy for the entire input at once instead of row by row.
Added new imports for predictors, data modules, and noise schedule.
Added new types and predictors for unconditional and sequence-to-sequence predictions.
Removed unused imports from the __init__.py file.
|
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #54 +/- ##
=======================================
Coverage 37.79% 37.79%
=======================================
Files 131 131
Lines 12416 12416
Branches 1711 1711
=======================================
Hits 4693 4693
Misses 7430 7430
Partials 293 293 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
3 participants
What
Adds BD3-LM (Block Discrete Denoising Diffusion Language Model) to
xlm-models,ported from the kuleshov-group reference implementation.
Contribution type
xlm-models/)Labels
model,enhancementChanges
xlm-models/bd3lm/packagetiny,small,medium+pretrained=autoTesting
Verified by running the eval configs: 1.000 exact match on
star-small and star-medium, at batch sizes 1, 4 and 16.
Fine-tuning from
kuleshov-group/bd3lm-owt-block_size4loads and trains.Documentation
xlm-models/bd3lm/README.mddocs/models/bd3lm.mdRelated issues