Skip to content

Converter for Llama based Masked Diffusion Models (Based on Dream) #263

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 10 commits into
base: main
Choose a base branch
from

Conversation

nitsanluke
Copy link
Contributor

@nitsanluke nitsanluke commented May 13, 2025

✨ Description

This PR provides a converter for Diffusion models based on Llama (and Dream). It complements the mask-diffusion training PR #238 and needs to merge after.

🔍 Type of change

Select all that apply:

  • 🐛 Bug fix (non-breaking change that addresses a specific issue)
  • 🚀 New feature (non-breaking change that adds functionality)
  • ⚠️ Breaking change (a change that could affect existing functionality)
  • 📈 Performance improvement/optimization (improves speed, memory usage, or efficiency)
  • 🛠️ Code refactor (non-functional changes that improve code readability, structure, etc.)
  • 📦 Dependency bump (updates dependencies, including Dockerfile or package changes)
  • 📝 Documentation change (updates documentation, including new content or typo fixes)
  • 🔧 Infrastructure/Build change (affects build process, CI/CD, or dependencies)

📝 Changes

List the key changes introduced in this PR:

  1. DiffusionLlama converter: For us to export our masked-LLM into a HF model
  2. Dream converter: For us to test Dream models

✅ Checklist

Make sure the following tasks are completed before submitting the PR:

General

  • 📜 I have read and followed the contributing guidelines.
  • 🏷️ I am using a clear and descriptive PR title that summarizes the key change or feature introduced.
  • 🎉 The functionality is complete, and I have tested the changes.
  • 📝 I have updated the documentation if needed.
  • ⚠️ The change does not introduce any new issues (e.g., runtime warnings, type checker errors, linting problems, unhandled edge cases).
  • 🧩 I have commented my code, especially in hard-to-understand areas.

Dependencies and Configuration

  • 🐋 I have updated the Docker configuration or dependencies, if applicable.
  • 🔄 I have ensured compatibility with the existing setup after dependency changes.

Testing

  • 🧪 I have added or updated tests to cover my changes.
  • ✔️ New and existing tests pass locally with my changes.
  • 🚦 I have tested these changes on GPUs and verified training stability.
  • 🏋️ I have tested the changes on realistic training workloads, if applicable.

Performance Impact

  • 📊 I have run benchmarks where applicable to evaluate the performance impact.
  • ✅ The benchmarks show no performance regression.
  • 🚀 The benchmarks indicate a potential performance improvement.
  • ⚠️ The benchmarks indicate a potential performance degradation.
  • 📈 I have provided benchmark results and detailed any performance impact below, if applicable.

📊 Performance Impact Details

If there is any impact on performance, describe it and provide benchmark results, if applicable:


🗒️ Additional Notes

tests/test_checkpoint.py ...............s.Fssss.

@nitsanluke nitsanluke changed the title WIP: Converter for Llama based Masked Diffusion Models (Based on Dream) Converter for Llama based Masked Diffusion Models (Based on Dream) May 14, 2025
@nitsanluke nitsanluke marked this pull request as ready for review May 14, 2025 17:57
@@ -133,6 +133,7 @@ class CustomModelingExportMixin:
modeling_file: typing.ClassVar[str]
configuration_file: typing.ClassVar[str]
configuration_cls: typing.ClassVar[type[PretrainedConfig]]
generation_utils_file: typing.ClassVar[typing.Optional[str]] = None
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

str|None

@@ -121,7 +121,7 @@ def test_convert_distributed_to_fast_llm():


@pytest.mark.depends(on=["test_convert_distributed_to_fast_llm"])
def test_convert_fast_llm_to_huggingface():
def test_convert_fast_llm_to_huggingface():
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please run the formatting hook.

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.

2 participants