Upload folder using huggingface_hub

ACKNOWLEDGMENTS

@@ -0,0 +1,31 @@
+Acknowledgements
+Portions of this ml-fs-dfm Software may utilize the following copyrighted
+material, the use of which is hereby acknowledged.
+
+_____________________
+
+Facebook, Inc. (Flow Matching)
+    Attribution-NonCommercial 4.0 International
+
+    Creative Commons Corporation ("Creative Commons") is not a law firm and
+    does not provide legal services or legal advice. Distribution of
+    Creative Commons public licenses does not create a lawyer-client or
+    other relationship. Creative Commons makes its licenses and related
+    information available on an "as-is" basis. Creative Commons gives no
+    warranties regarding its licenses, any material licensed under their
+    terms and conditions, or any related information. Creative Commons
+    disclaims all liability for damages resulting from their use to the
+    fullest extent possible.
+
+    By exercising the Licensed Rights (defined below), You accept and agree
+    to be bound by the terms and conditions of this Creative Commons
+    Attribution-NonCommercial 4.0 International Public License ("Public
+    License"). To the extent this Public License may be interpreted as a
+    contract, You are granted the Licensed Rights in consideration of Your
+    acceptance of these terms and conditions, and the Licensor grants You
+    such rights in consideration of benefits the Licensor receives from
+    making the Licensed Material available under these terms and
+    conditions.
+
+    For the full license text, see:
+    https://creativecommons.org/licenses/by-nc/4.0/legalcode

LICENSE

@@ -0,0 +1,41 @@
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+
+IMPORTANT:  This Apple software is supplied to you by Apple
+Inc. ("Apple") in consideration of your agreement to the following
+terms, and your use, installation, modification or redistribution of
+this Apple software constitutes acceptance of these terms.  If you do
+not agree with these terms, please do not use, install, modify or
+redistribute this Apple software.
+
+In consideration of your agreement to abide by the following terms, and
+subject to these terms, Apple grants you a personal, non-exclusive
+license, under Apple's copyrights in this original Apple software (the
+"Apple Software"), to use, reproduce, modify and redistribute the Apple
+Software, with or without modifications, in source and/or binary forms;
+provided that if you redistribute the Apple Software in its entirety and
+without modifications, you must retain this notice and the following
+text and disclaimers in all such redistributions of the Apple Software.
+Neither the name, trademarks, service marks or logos of Apple Inc. may
+be used to endorse or promote products derived from the Apple Software
+without specific prior written permission from Apple.  Except as
+expressly stated in this notice, no other rights or licenses, express or
+implied, are granted by Apple herein, including but not limited to any
+patent rights that may be infringed by your derivative works or by other
+works in which the Apple Software may be incorporated.
+
+The Apple Software is provided by Apple on an "AS IS" basis.  APPLE
+MAKES NO WARRANTIES, EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION
+THE IMPLIED WARRANTIES OF NON-INFRINGEMENT, MERCHANTABILITY AND FITNESS
+FOR A PARTICULAR PURPOSE, REGARDING THE APPLE SOFTWARE OR ITS USE AND
+OPERATION ALONE OR IN COMBINATION WITH YOUR PRODUCTS.
+
+IN NO EVENT SHALL APPLE BE LIABLE FOR ANY SPECIAL, INDIRECT, INCIDENTAL
+OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+INTERRUPTION) ARISING IN ANY WAY OUT OF THE USE, REPRODUCTION,
+MODIFICATION AND/OR DISTRIBUTION OF THE APPLE SOFTWARE, HOWEVER CAUSED
+AND WHETHER UNDER THEORY OF CONTRACT, TORT (INCLUDING NEGLIGENCE),
+STRICT LIABILITY OR OTHERWISE, EVEN IF APPLE HAS BEEN ADVISED OF THE
+POSSIBILITY OF SUCH DAMAGE.
+
+Third-party software acknowledgments are contained in the file named ACKNOWLEDGMENTS.

README.md

@@ -0,0 +1,54 @@
+---
+license: other
+tags:
+  - non-commercial
+  - text-generation
+  - flow-matching
+datasets:
+  - cerebras/SlimPajama-627B
+---
+
+# DFM
+
+## Summary
+`DFM` is a continued-pretraining checkpoint based on Apple's fs-dfm weights. It is trained with Flow Matching code and released for research/non-commercial use only.
+
+Base checkpoint (external, not on HF):
+```
+https://ml-site.cdn-apple.com/models/fs-dfm/checkpoint.pth
+```
+
+## Training
+- Continued pretraining from Apple's fs-dfm checkpoint
+- Dataset: SlimPajama-627B
+- Steps: 250,000
+- Global batch size: 256
+
+## License
+Research/non-commercial use only. This repository is governed by the Apple Software License (see `LICENSE`) and includes non-commercial restrictions inherited from Flow Matching (CC BY-NC 4.0). See `ACKNOWLEDGMENTS` for third-party notices.
+
+## Intended Use
+Research and non-commercial use only.
+
+## Limitations
+Commercial use is not permitted. Dataset-specific licensing constraints apply to SlimPajama's underlying sources.
+
+## Usage
+### Hugging Face (trust_remote_code)
+This repo provides `configuration_dfm.py` and `modeling_dfm.py` for HF loading with `trust_remote_code=True`.
+
+Example:
+```python
+from transformers import AutoConfig, AutoModel
+
+config = AutoConfig.from_pretrained(".", trust_remote_code=True)
+model = AutoModel.from_pretrained(".", trust_remote_code=True)
+```
+
+Note:
+- This model expects `x_t` and `time` inputs (flow-matching style), not GPT-style autoregressive inputs.
+
+This release includes model-only weights (`model.safetensors`) for inference/forward passes. Full training/eval/sampling code is available in the original project: `https://github.com/apple/ml-fs-dfm`.
+
+## Acknowledgments
+This model is derived from Apple's fs-dfm checkpoint and follows the original Apple license terms. The original project is at `https://github.com/apple/ml-fs-dfm`. See `ACKNOWLEDGMENTS` for third-party attributions and licensing.

config.json

@@ -0,0 +1,31 @@
+{
+  "model_type": "dfm",
+  "architectures": [
+    "DFMModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_dfm.DFMConfig",
+    "AutoModel": "modeling_dfm.DFMModel"
+  },
+  "vocab_size": 50257,
+  "hidden_size": 2048,
+  "cond_dim": 256,
+  "num_hidden_layers": 21,
+  "n_blocks": 21,
+  "num_attention_heads": 32,
+  "n_heads": 32,
+  "max_position_embeddings": 1024,
+  "sequence_length": 1024,
+  "dropout": 0.1,
+  "rotary_dim": 64,
+  "source_distribution": "mask",
+  "flow_scheduler_type": "polynomial",
+  "flow_exponent": 1.0,
+  "flow_loss_function": "generalized_kl",
+  "sampling_steps": 1024,
+  "bos_token_id": 50256,
+  "eos_token_id": 50256,
+  "mask_token_id": 50257,
+  "tokenizer_name": "gpt2",
+  "dtype": "bfloat16"
+}

configuration_dfm.py

@@ -0,0 +1,47 @@
+from transformers import PretrainedConfig
+
+
+class DFMConfig(PretrainedConfig):
+    model_type = "dfm"
+
+    def __init__(
+        self,
+        vocab_size=50257,
+        hidden_size=2048,
+        cond_dim=256,
+        n_blocks=21,
+        n_heads=32,
+        dropout=0.1,
+        sequence_length=1024,
+        source_distribution="mask",
+        flow_scheduler_type="polynomial",
+        flow_exponent=1.0,
+        flow_loss_function="generalized_kl",
+        sampling_steps=1024,
+        bos_token_id=50256,
+        eos_token_id=50256,
+        mask_token_id=50257,
+        tokenizer_name="gpt2",
+        dtype="bfloat16",
+        **kwargs,
+    ):
+        super().__init__(
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            **kwargs,
+        )
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.cond_dim = cond_dim
+        self.n_blocks = n_blocks
+        self.n_heads = n_heads
+        self.dropout = dropout
+        self.sequence_length = sequence_length
+        self.source_distribution = source_distribution
+        self.flow_scheduler_type = flow_scheduler_type
+        self.flow_exponent = flow_exponent
+        self.flow_loss_function = flow_loss_function
+        self.sampling_steps = sampling_steps
+        self.mask_token_id = mask_token_id
+        self.tokenizer_name = tokenizer_name
+        self.dtype = dtype

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d80bdb27307852691fa4eb6eddd880307cb07a12cd30458bc051c8ff23662291
+size 5322735224

modeling_dfm.py

@@ -0,0 +1,372 @@
+import math
+from types import SimpleNamespace
+from typing import Optional, Tuple
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange, repeat
+from torch import Tensor, nn
+from transformers import PreTrainedModel
+
+try:
+    import flash_attn
+except ImportError:
+    flash_attn = None
+
+try:
+    import flash_attn_interface
+except ImportError:
+    flash_attn_interface = None
+from configuration_dfm import DFMConfig
+
+
+class Rotary(torch.nn.Module):
+    """
+    From: https://github.com/louaaron/Score-Entropy-Discrete-Diffusion
+    """
+
+    def __init__(self, dim: int, base: int = 10_000):
+        super().__init__()
+        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer("inv_freq", inv_freq)
+        self.seq_len_cached = None
+        self.cos_cached = None
+        self.sin_cached = None
+
+    def forward(self, x: Tensor, seq_dim: int = 1) -> Tuple[Tensor, Tensor]:
+        seq_len = x.shape[seq_dim]
+        if seq_len != self.seq_len_cached:
+            self.seq_len_cached = seq_len
+            t = torch.arange(x.shape[seq_dim], device=x.device).type_as(self.inv_freq)
+            freqs = torch.einsum("i,j->ij", t, self.inv_freq.clone())
+            emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
+
+            # dims are: batch, seq_len, qkv, head, dim
+            self.cos_cached = emb.cos()[None, :, None, None, :].repeat(1, 1, 3, 1, 1)
+            self.sin_cached = emb.sin()[None, :, None, None, :].repeat(1, 1, 3, 1, 1)
+
+            # This makes the transformation on v an identity.
+            self.cos_cached[:, :, 2, :, :].fill_(1.0)
+            self.sin_cached[:, :, 2, :, :].fill_(0.0)
+
+        return self.cos_cached, self.sin_cached
+
+
+def rotate_half(x: Tensor) -> Tensor:
+    x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :]
+
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_emb_torch(x, cos, sin, interleaved=False):
+    """
+    From: https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/layers/rotary.py#L20
+    """
+    cos = cos[0, :, 0, 0, : cos.shape[-1] // 2]
+    sin = sin[0, :, 0, 0, : sin.shape[-1] // 2]
+
+    ro_dim = cos.shape[-1] * 2
+    assert ro_dim <= x.shape[-1]
+    cos = repeat(
+        cos, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)"
+    )
+    sin = repeat(
+        sin, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)"
+    )
+
+    return x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim]) * sin
+
+
+def bias_dropout_add_scale(
+    x: Tensor, scale: Tensor, residual: Optional[Tensor], prob: float, training: bool
+) -> Tensor:
+    return residual + scale * F.dropout(x, p=prob, training=training)
+
+
+def modulate(x: Tensor, shift: Tensor, scale: Tensor) -> Tensor:
+    return x * (1 + scale) + shift
+
+
+class LayerNorm(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones([dim]))
+        self.dim = dim
+
+    def forward(self, x: Tensor) -> Tensor:
+        with torch.amp.autocast("cuda", enabled=False):
+            x = F.layer_norm(x.float(), [self.dim])
+
+        return x * self.weight[None, None, :]
+
+
+class TimestepEmbedder(nn.Module):
+    """
+    Embeds scalar timesteps into vector representations.
+    """
+
+    def __init__(self, hidden_size: int, frequency_embedding_size: int = 256):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+
+    @staticmethod
+    def timestep_embedding(time: Tensor, dim: int, max_period: int = 10000) -> Tensor:
+        """
+        Create sinusoidal timestep embeddings.
+        :param t: a 1-D Tensor of N indices, one per batch element.
+                          These may be fractional.
+        :param dim: the dimension of the output.
+        :param max_period: controls the minimum frequency of the embeddings.
+        :return: an (N, D) Tensor of positional embeddings.
+        """
+        half = dim // 2
+        freqs = torch.exp(
+            -math.log(max_period)
+            * torch.arange(start=0, end=half, dtype=torch.float32)
+            / half
+        ).to(device=time.device)
+        args = time[:, None].float() * freqs[None]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat(
+                [embedding, torch.zeros_like(embedding[:, :1])], dim=-1
+            )
+        return embedding
+
+    def forward(self, time: Tensor) -> Tensor:
+        t_freq = self.timestep_embedding(time=time, dim=self.frequency_embedding_size)
+        t_emb = self.mlp(t_freq)
+        return t_emb
+
+
+class DDiTBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        n_heads: int,
+        cond_dim: int,
+        mlp_ratio: int = 4,
+        dropout: float = 0.1,
+    ):
+        super().__init__()
+        assert dim % n_heads == 0, "dim must be devisable by n_heads"
+
+        self.n_heads = n_heads
+        self.dim = dim
+        self.dropout = dropout
+
+        self.head_dim = self.dim // self.n_heads
+
+        self.norm1 = LayerNorm(dim=dim)
+
+        self.qw = nn.Linear(dim, dim, bias=False)
+        self.kw = nn.Linear(dim, dim, bias=False)
+        self.vw = nn.Linear(dim, dim, bias=False)
+
+        self.attn_out = nn.Linear(dim, dim, bias=False)
+        self.dropout1 = nn.Dropout(dropout)
+
+        self.norm2 = LayerNorm(dim=dim)
+        self.mlp = nn.Sequential(
+            nn.Linear(dim, mlp_ratio * dim, bias=True),
+            nn.GELU(approximate="tanh"),
+            nn.Linear(mlp_ratio * dim, dim, bias=True),
+        )
+
+        self.adaLN_modulation = nn.Linear(cond_dim, 6 * dim, bias=True)
+        self.adaLN_modulation.weight.data.zero_()
+        self.adaLN_modulation.bias.data.zero_()
+
+    def forward(self, x: Tensor, rotary_cos_sin: Tensor, c: Tensor) -> Tensor:
+        batch_size, seq_len = x.shape[0], x.shape[1]
+
+        (
+            shift_msa,
+            scale_msa,
+            gate_msa,
+            shift_mlp,
+            scale_mlp,
+            gate_mlp,
+        ) = self.adaLN_modulation(c)[:, None].chunk(6, dim=2)
+
+        x_skip = x
+        x = modulate(x=self.norm1(x), shift=shift_msa, scale=scale_msa)
+
+        q = self.qw(x)
+        k = self.kw(x)
+        v = self.vw(x)
+
+        q, k, v = (
+            item.view(batch_size, seq_len, self.n_heads, self.head_dim)
+            for item in (q, k, v)
+        )
+
+        with torch.amp.autocast("cuda", enabled=False):
+            cos, sin = rotary_cos_sin
+            original_dtype = q.dtype
+
+            q = apply_rotary_emb_torch(
+                x=q.float(), cos=cos.float(), sin=sin.float()
+            ).to(original_dtype)
+            k = apply_rotary_emb_torch(
+                x=k.float(), cos=cos.float(), sin=sin.float()
+            ).to(original_dtype)
+
+        use_flash_attn = (
+            flash_attn_interface is not None or flash_attn is not None
+        ) and q.is_cuda
+        if use_flash_attn:
+            qkv = torch.stack((q, k, v), dim=2)
+            if flash_attn_interface is not None:
+                x = flash_attn_interface.flash_attn_qkvpacked_func(qkv, causal=False)
+            else:
+                x = flash_attn.flash_attn_qkvpacked_func(qkv, 0.0, causal=False)
+            x = rearrange(x, "b s h d -> b s (h d)", b=batch_size)
+        else:
+            q, k, v = (item.transpose(1, 2) for item in (q, k, v))
+            x = F.scaled_dot_product_attention(query=q, key=k, value=v)
+            x = rearrange(x, "b h s d -> b s (h d)", b=batch_size)
+        x = bias_dropout_add_scale(
+            x=self.attn_out(x),
+            scale=gate_msa,
+            residual=x_skip,
+            prob=self.dropout,
+            training=self.training,
+        )
+        x = bias_dropout_add_scale(
+            x=self.mlp(modulate(x=self.norm2(x), shift=shift_mlp, scale=scale_mlp)),
+            scale=gate_mlp,
+            residual=x,
+            prob=self.dropout,
+            training=self.training,
+        )
+
+        return x
+
+
+class DDitFinalLayer(nn.Module):
+    def __init__(self, hidden_size: int, out_channels: int, cond_dim: int):
+        super().__init__()
+        self.norm_final = LayerNorm(hidden_size)
+        self.linear = nn.Linear(hidden_size, out_channels)
+        self.linear.weight.data.zero_()
+        self.linear.bias.data.zero_()
+
+        self.adaLN_modulation = nn.Linear(cond_dim, 2 * hidden_size, bias=True)
+        self.adaLN_modulation.weight.data.zero_()
+        self.adaLN_modulation.bias.data.zero_()
+
+    def forward(self, x: Tensor, c: Tensor) -> Tensor:
+        shift, scale = self.adaLN_modulation(c)[:, None].chunk(2, dim=2)
+        x = modulate(x=self.norm_final(x), shift=shift, scale=scale)
+        x = self.linear(x)
+
+        return x
+
+
+class Transformer(nn.Module):
+    def __init__(self, vocab_size: int, masked: bool, config):
+        super().__init__()
+
+        if isinstance(config, dict):
+            config = SimpleNamespace(**config)
+
+        self.config = config
+        self.vocab_size = vocab_size
+
+        add_token = 1 if masked else 0
+
+        self.vocab_embed = nn.Embedding(self.vocab_size + add_token, config.hidden_size)
+
+        self.time_embedding = TimestepEmbedder(hidden_size=config.cond_dim)
+        self.rotary_emb = Rotary(dim=config.hidden_size // config.n_heads)
+
+        self.blocks = nn.ModuleList(
+            [
+                DDiTBlock(
+                    dim=config.hidden_size,
+                    n_heads=config.n_heads,
+                    cond_dim=config.cond_dim,
+                    dropout=config.dropout,
+                )
+                for _ in range(config.n_blocks)
+            ]
+        )
+
+        self.output_layer = DDitFinalLayer(
+            hidden_size=config.hidden_size,
+            out_channels=vocab_size + add_token,
+            cond_dim=config.cond_dim,
+        )
+
+    def forward(self, x_t: Tensor, time: Tensor) -> Tensor:
+        x = self.vocab_embed(x_t)
+        c = F.silu(self.time_embedding(time=time))
+
+        rotary_cos_sin = self.rotary_emb(x=x)
+
+        with torch.amp.autocast("cuda", dtype=torch.bfloat16):
+            for i in range(len(self.blocks)):
+                x = self.blocks[i](x=x, rotary_cos_sin=rotary_cos_sin, c=c)
+
+            x = self.output_layer(x=x, c=c)
+
+        return x
+
+
+class DFMModel(PreTrainedModel):
+    config_class = DFMConfig
+    base_model_prefix = "model"
+
+    def __init__(self, config: DFMConfig):
+        super().__init__(config)
+        masked = config.source_distribution == "mask"
+        self.model = Transformer(
+            vocab_size=config.vocab_size,
+            masked=masked,
+            config={
+                "hidden_size": config.hidden_size,
+                "cond_dim": config.cond_dim,
+                "length": config.sequence_length,
+                "n_blocks": config.n_blocks,
+                "n_heads": config.n_heads,
+                "dropout": config.dropout,
+                "compile": False,
+            },
+        )
+        self.post_init()
+
+    def forward(
+        self,
+        x_t: torch.Tensor,
+        time: torch.Tensor,
+        **kwargs,
+    ) -> torch.Tensor:
+        return self.model(x_t=x_t, time=time)
+
+    @classmethod
+    def _load_pretrained_model(
+        cls,
+        model,
+        state_dict,
+        *args,
+        **kwargs,
+    ):
+        if state_dict is not None:
+            if "model" in state_dict and isinstance(state_dict["model"], dict):
+                state_dict = state_dict["model"]
+            if state_dict and not any(
+                k.startswith("model.") for k in state_dict.keys()
+            ):
+                state_dict = {f"model.{k}": v for k, v in state_dict.items()}
+        return super()._load_pretrained_model(
+            model,
+            state_dict,
+            *args,
+            **kwargs,
+        )

special_tokens_map.json

@@ -0,0 +1,5 @@
+{
+  "bos_token": "<|endoftext|>",
+  "eos_token": "<|endoftext|>",
+  "unk_token": "<|endoftext|>"
+}

tokenizer_config.json

@@ -0,0 +1,20 @@
+{
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "50256": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<|endoftext|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "extra_special_tokens": {},
+  "model_max_length": 1024,
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|endoftext|>"
+}