Skip to content

Add forward-only LLaDA architecture adapter#1512

Merged
jlarson4 merged 2 commits into
TransformerLensOrg:devfrom
hsusul:llada-adapter
Jul 14, 2026
Merged

Add forward-only LLaDA architecture adapter#1512
jlarson4 merged 2 commits into
TransformerLensOrg:devfrom
hsusul:llada-adapter

Conversation

@hsusul

@hsusul hsusul commented Jul 12, 2026

Copy link
Copy Markdown
Contributor

Summary

Addresses #1465.

Add TransformerBridge support for LLaDAModelLM single-pass inference.

  • Map LLaDA residual, attention, MLP, normalization, and logits paths.
  • Preserve bidirectional attention and padding-mask behavior.
  • Expose standard TransformerLens hook points.
  • Require explicit trust_remote_code opt-in.
  • Reject unsupported causal generation and shifted-token loss.
  • Add offline numerical-equivalence and framework regression tests.

Scope

This PR supports a single masked-token forward pass through TransformerBridge.

Iterative diffusion generation, remasking, confidence schedules, and diffusion-timestep hooks remain out of scope.

Validation

  • Offline tiny-model numerical equivalence
  • Bidirectional and padding-mask behavior
  • Intermediate hook equivalence
  • Full state-dict mapping coverage
  • FP32 and BF16 coverage
  • Existing causal-model regressions
  • make unit-test: 3,065 passed

@jlarson4 jlarson4 changed the base branch from main to dev July 14, 2026 14:37
@jlarson4

Copy link
Copy Markdown
Collaborator

@hsusul Great work on this! Looks good, the limitations seem reasonable for initial support. I merged the latest dev into your branch to resolve any conflicts, and will merge once the CI passes

@jlarson4 jlarson4 linked an issue Jul 14, 2026 that may be closed by this pull request
1 task
@jlarson4 jlarson4 merged commit 983cb38 into TransformerLensOrg:dev Jul 14, 2026
25 checks passed
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.

[Proposal] Add LLaDA masked-diffusion adapter (LLaDAModelLM)

2 participants