Add files using upload-large-folder tool

checkpoints/official/duo_dcd/owt/model.py

@@ -0,0 +1,601 @@
+import math
+import typing
+
+import einops
+import flash_attn
+import flash_attn.layers.rotary
+import huggingface_hub
+import omegaconf
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import transformers
+
+from .config import DUOConfig
+
+# Flags required to enable jit fusion kernels
+torch._C._jit_set_profiling_mode(False)
+torch._C._jit_set_profiling_executor(False)
+torch._C._jit_override_can_fuse_on_cpu(True)
+torch._C._jit_override_can_fuse_on_gpu(True)
+
+
+def bias_dropout_add_scale(
+    x: torch.Tensor,
+    bias: typing.Optional[torch.Tensor],
+    scale: torch.Tensor,
+    residual: typing.Optional[torch.Tensor],
+    prob: float,
+    training: bool) -> torch.Tensor:
+  if bias is not None:
+    out = scale * F.dropout(x + bias, p=prob, training=training)
+  else:
+    out = scale * F.dropout(x, p=prob, training=training)
+
+  if residual is not None:
+    out = residual + out
+  return out
+
+
+def get_bias_dropout_add_scale(training):
+  def _bias_dropout_add(x, bias, scale, residual, prob):
+    return bias_dropout_add_scale(
+      x, bias, scale, residual, prob, training)
+
+  return _bias_dropout_add
+
+
+# function overload
+def modulate(x: torch.Tensor,
+             shift: torch.Tensor,
+             scale: torch.Tensor) -> torch.Tensor:
+  return x * (1 + scale) + shift
+
+
+@torch.jit.script
+def bias_dropout_add_scale_fused_train(
+    x: torch.Tensor,
+    bias: typing.Optional[torch.Tensor],
+    scale: torch.Tensor,
+    residual: typing.Optional[torch.Tensor],
+    prob: float) -> torch.Tensor:
+  return bias_dropout_add_scale(
+    x, bias, scale, residual, prob, True)
+
+
+@torch.jit.script
+def bias_dropout_add_scale_fused_inference(
+    x: torch.Tensor,
+    bias: typing.Optional[torch.Tensor],
+    scale: torch.Tensor,
+    residual: typing.Optional[torch.Tensor],
+    prob: float) -> torch.Tensor:
+  return bias_dropout_add_scale(
+    x, bias, scale, residual, prob, False)
+
+
+@torch.jit.script
+def modulate_fused(x: torch.Tensor,
+                   shift: torch.Tensor,
+                   scale: torch.Tensor) -> torch.Tensor:
+  return modulate(x, shift, scale)
+
+
+class Rotary(torch.nn.Module):
+  def __init__(self, dim, base=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, seq_dim=1):
+    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.)
+      self.sin_cached[:,:,2,:,:].fill_(0.)
+
+    return self.cos_cached, self.sin_cached
+
+
+def rotate_half(x):
+  x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :]
+  return torch.cat((-x2, x1), dim=-1)
+
+
+def split_and_apply_rotary_pos_emb(qkv, rotary_cos_sin):
+  with torch.cuda.amp.autocast(enabled=False):
+    cos, sin = rotary_cos_sin
+    cos = cos.to(qkv.dtype)
+    sin = sin.to(qkv.dtype)
+    cos = cos[0,:,0,0,:cos.shape[-1]//2]
+    sin = sin[0,:,0,0,:sin.shape[-1]//2]
+    q, k, v = qkv.chunk(3, dim=2)
+    q = flash_attn.layers.rotary.apply_rotary_emb_torch(
+      q.squeeze(dim=2), cos, sin)
+    k = flash_attn.layers.rotary.apply_rotary_emb_torch(
+      k.squeeze(dim=2), cos, sin)
+    v = v.squeeze(dim=2)
+  return q, k, v
+
+
+def apply_rotary_pos_emb(qkv, cos, sin):
+  cos = cos[0,:,0,0,:cos.shape[-1]//2]
+  sin = sin[0,:,0,0,:sin.shape[-1]//2]
+  return flash_attn.layers.rotary.apply_rotary_emb_qkv_(qkv, cos, sin)
+
+
+def regular_attention_multi_headed(q, k, v):
+  # Assuming qkv is a tensor with shape [batch, seq_len, 3, num_heads, head_dim]
+  # where the 3 represents Q, K, V packed in that order
+  attention_output = F.scaled_dot_product_attention(
+    query=q.transpose(1, 2),
+    key=k.transpose(1, 2),
+    value=v.transpose(1, 2),
+    attn_mask=None,
+    dropout_p=0.0,
+    is_causal=False)
+  # [batch_size, seq_len, num_heads, head_dim]
+  attention_output = attention_output.transpose(1, 2)
+  return einops.rearrange(attention_output, 'b s h d -> b s (h d)')
+
+
+#################################################################################
+#                                  Layers                                       #
+#################################################################################
+class LayerNorm(nn.Module):
+  def __init__(self, dim):
+    super().__init__()
+    self.weight = nn.Parameter(torch.ones([dim]))
+    self.dim = dim
+  def forward(self, x):
+    with torch.cuda.amp.autocast(enabled=False):
+      x = F.layer_norm(x.float(), [self.dim])
+    return x * self.weight[None, None, :]
+
+
+def residual_linear(x, W, x_skip, residual_scale):
+  """x_skip + residual_scale * W @ x"""
+  dim_out, dim_in = W.shape[0], W.shape[1]
+  return torch.addmm(
+    x_skip.view(-1, dim_out),
+    x.view(-1, dim_in),
+    W.T,
+    alpha=residual_scale).view(*x.shape[:-1], dim_out)
+
+
+#################################################################################
+#               Embedding Layers for Timesteps and Class Labels                 #
+#################################################################################
+class TimestepEmbedder(nn.Module):
+  """
+  Embeds scalar timesteps into vector representations.
+  """
+  def __init__(self, hidden_size, frequency_embedding_size=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(t, dim, max_period=10000):
+    """
+    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.
+    """
+    # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
+    half = dim // 2
+    freqs = torch.exp(
+      - math.log(max_period)
+      * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device)
+      / half)
+    args = t[:, 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, t):
+    t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+    t_emb = self.mlp(t_freq)
+    return t_emb
+
+
+class LabelEmbedder(nn.Module):
+  """Embeds class labels into vector representations.
+  
+  Also handles label dropout for classifier-free guidance.
+  """
+  def __init__(self, num_classes, cond_size):
+    super().__init__()
+    self.embedding_table = nn.Embedding(num_classes + 1, cond_size)
+    self.num_classes = num_classes
+
+    # TODO think of initializing with 0.02 std deviation like in original DiT paper
+
+  def forward(self, labels):
+    embeddings = self.embedding_table(labels)
+    return embeddings
+    
+
+#################################################################################
+#                                 Core Model                                    #
+#################################################################################
+
+class DDiTBlockCausal(nn.Module):
+  def __init__(self, dim, n_heads, mlp_ratio=4, dropout=0.1):
+    super().__init__()
+    self.n_heads = n_heads
+
+    self.norm1 = LayerNorm(dim)
+    self.attn_qkv = nn.Linear(dim, 3 * dim, bias=False)
+    self.attn_out = nn.Linear(dim, dim, bias=False)
+    self.dropout1 = nn.Dropout(dropout)
+
+    self.norm2 = LayerNorm(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.dropout2 = nn.Dropout(dropout)
+    self.dropout = dropout
+
+  def _get_bias_dropout_scale(self):
+    if self.training:
+      return bias_dropout_add_scale_fused_train
+    else:
+      return bias_dropout_add_scale_fused_inference
+
+  def forward(self, x, rotary_cos_sin, **kwargs):
+    del kwargs
+    batch_size, seq_len = x.shape[0], x.shape[1]
+
+    bias_dropout_scale_fn = self._get_bias_dropout_scale()
+
+    # attention operation
+    x_skip = x
+    x = self.norm1(x)
+
+    qkv = self.attn_qkv(x)
+    qkv = einops.rearrange(
+      qkv,
+      'b s (three h d) -> b s three h d',
+      three=3,
+      h=self.n_heads)
+    with torch.cuda.amp.autocast(enabled=False):
+      cos, sin = rotary_cos_sin
+      qkv = apply_rotary_pos_emb(
+        qkv, cos.to(qkv.dtype), sin.to(qkv.dtype)
+      )
+    qkv = einops.rearrange(qkv, 'b s ... -> (b s) ...')
+    cu_seqlens = torch.arange(
+      0, (batch_size + 1) * seq_len,
+      step=seq_len, dtype=torch.int32, device=qkv.device)
+    x = flash_attn.flash_attn_interface.flash_attn_varlen_qkvpacked_func(
+      qkv, cu_seqlens, seq_len, 0.0, causal=True)
+
+    x = einops.rearrange(x, '(b s) h d -> b s (h d)',
+                         b=batch_size)
+
+    scale = torch.ones(1, device=x.device, dtype=x.dtype)
+    x = bias_dropout_scale_fn(
+      self.attn_out(x), None, scale, x_skip, self.dropout)
+
+    # mlp operation
+    x = bias_dropout_scale_fn(
+      self.mlp(self.norm2(x)), None, scale, x, self.dropout)
+    return x
+
+
+
+class DDiTBlock(nn.Module):
+  def __init__(self, dim, n_heads, adaLN,
+               cond_dim=None, mlp_ratio=4,
+               dropout=0.1):
+    super().__init__()
+    self.n_heads = n_heads
+    self.adaLN = adaLN
+
+    self.norm1 = LayerNorm(dim)
+    self.attn_qkv = nn.Linear(dim, 3 * dim, bias=False)
+    self.attn_out = nn.Linear(dim, dim, bias=False)
+    self.dropout1 = nn.Dropout(dropout)
+
+    self.norm2 = LayerNorm(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.dropout2 = nn.Dropout(dropout)
+    self.dropout = dropout
+
+    if self.adaLN:
+      self.adaLN_modulation = nn.Linear(cond_dim, 6 * dim)
+      self.adaLN_modulation.weight.data.zero_()
+      self.adaLN_modulation.bias.data.zero_()
+
+
+  def _get_bias_dropout_scale(self):
+    if self.training:
+      return bias_dropout_add_scale_fused_train
+    else:
+      return bias_dropout_add_scale_fused_inference
+
+
+  def forward(self, x, rotary_cos_sin, c=None):
+
+    bias_dropout_scale_fn = self._get_bias_dropout_scale()
+
+    x_skip = x
+    x = self.norm1(x)
+
+    if self.adaLN:
+      # self.adaLN_modulation(c): (128, 1536)
+      # self.adaLN_modulation(c)[:, None]: (128, 1, 1536)
+      # "" .chunk(6, dim=2) returns 6 tuples of shapes (128, 1, 256)
+      (shift_msa, scale_msa, gate_msa, shift_mlp,
+       scale_mlp, gate_mlp) = self.adaLN_modulation(c)[:, None].chunk(6, dim=2)
+      x = modulate_fused(x, shift_msa, scale_msa)
+
+    qkv = einops.rearrange(
+      self.attn_qkv(x),
+      'b s (three h d) -> b s three h d',
+      three=3,
+      h=self.n_heads)
+    q, k, v = split_and_apply_rotary_pos_emb(qkv, rotary_cos_sin)
+    
+    x = regular_attention_multi_headed(q, k, v)
+
+    if self.adaLN:
+      x = bias_dropout_scale_fn(self.attn_out(x),
+                                None,
+                                gate_msa,
+                                x_skip,
+                                self.dropout)
+      x = bias_dropout_scale_fn(
+        self.mlp(modulate_fused(
+          self.norm2(x), shift_mlp, scale_mlp)),
+        None, gate_mlp, x, self.dropout)
+    else:
+      scale = torch.ones(1, device=x.device, dtype=x.dtype)
+      x = bias_dropout_scale_fn(
+        self.attn_out(x), None, scale, x_skip, self.dropout)
+      x = bias_dropout_scale_fn(
+        self.mlp(self.norm2(x)), None, scale, x, self.dropout)
+    return x
+
+
+class EmbeddingLayer(nn.Module):
+  def __init__(self, dim, vocab_dim):
+    super().__init__()
+    self.embedding = nn.Parameter(torch.empty((vocab_dim, dim)))
+    torch.nn.init.kaiming_uniform_(self.embedding, a=math.sqrt(5))
+
+  def forward(self, x):
+    if x.ndim == 2:
+      return self.embedding[x]
+    assert x.ndim == 3
+    return torch.einsum(
+      "blv,ve->ble",
+      torch.nn.functional.softmax(x, dim=-1).float(),
+      self.embedding.float()).to(x.dtype)
+
+
+class DDiTFinalLayer(nn.Module):
+  def __init__(self, hidden_size, out_channels, cond_dim,
+               adaLN):
+    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 = adaLN
+    if self.adaLN:
+      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, c):
+    x = self.norm_final(x)
+    if self.adaLN:
+      shift, scale = self.adaLN_modulation(c)[:, None].chunk(2, dim=2)
+      x = modulate_fused(x, shift, scale)
+    x = self.linear(x)
+    return x
+
+
+class DIT(nn.Module, huggingface_hub.PyTorchModelHubMixin):
+  def __init__(self, config, vocab_size: int):
+    super().__init__()
+    if type(config) == dict:
+      config = omegaconf.OmegaConf.create(config)
+    self.causal = config.algo.causal_attention
+    self.adaLN = not self.causal
+    self.config = config
+    self.vocab_size = vocab_size
+    dim = config.model.hidden_size
+    cond_dim = config.model.cond_dim
+    self.vocab_embed = EmbeddingLayer(dim, vocab_size)
+    if not self.causal:
+      self.sigma_map = TimestepEmbedder(cond_dim)
+    self.rotary_emb = Rotary(dim // config.model.n_heads)
+
+    blocks = []
+    for _ in range(config.model.n_blocks):
+      if self.causal:
+        block = DDiTBlockCausal(
+          dim=dim,
+          n_heads=config.model.n_heads,
+          dropout=config.model.dropout)
+      else:
+        block = DDiTBlock(
+          dim=dim,
+          n_heads=config.model.n_heads,
+          cond_dim=cond_dim,
+          adaLN=self.adaLN,
+          dropout=config.model.dropout)
+      blocks.append(block)
+    self.blocks = nn.ModuleList(blocks)
+
+    self.output_layer = DDiTFinalLayer(
+      hidden_size=dim,
+      out_channels=vocab_size,
+      cond_dim=cond_dim,
+      adaLN=self.adaLN)
+    self.scale_by_sigma = config.model.scale_by_sigma
+
+  def _get_bias_dropout_scale(self):
+    if self.training:
+      return bias_dropout_add_scale_fused_train
+    else:
+      return  bias_dropout_add_scale_fused_inference
+
+  def forward(self, x, sigma):
+    x = self.vocab_embed(x)
+    if self.causal:
+      t_cond = None
+    else:
+      t_cond = F.silu(self.sigma_map(sigma))
+
+    rotary_cos_sin = self.rotary_emb(x).to(x.device)
+
+    with torch.cuda.amp.autocast(dtype=torch.bfloat16):
+      for i in range(len(self.blocks)):
+        x = self.blocks[i](x, rotary_cos_sin, c=t_cond)
+      x = self.output_layer(x, c=t_cond)
+
+    return x
+
+
+
+class HFDIT(torch.nn.Module):
+  def __init__(self, config):
+    super().__init__()
+    self.causal = config.causal
+    self.adaLN = not self.causal
+    self.vocab_size = config.vocab_size
+    dim = config.hidden_dim
+    cond_dim = config.cond_dim
+    self.vocab_embed = EmbeddingLayer(dim, self.vocab_size)
+    if not self.causal:
+      self.sigma_map = TimestepEmbedder(cond_dim)
+    self.rotary_emb = Rotary(dim // config.n_heads)
+
+    blocks = []
+    for _ in range(config.n_blocks):
+      if self.causal:
+        block = DDiTBlockCausal(
+          dim=dim,
+          n_heads=config.n_heads,
+          dropout=config.dropout)
+      else:
+        block = DDiTBlock(
+          dim=dim,
+          n_heads=config.n_heads,
+          cond_dim=cond_dim,
+          adaLN=self.adaLN,
+          dropout=config.dropout)
+      blocks.append(block)
+    self.blocks = torch.nn.ModuleList(blocks)
+
+    self.output_layer = DDiTFinalLayer(
+      hidden_size=dim,
+      out_channels=self.vocab_size,
+      cond_dim=cond_dim,
+      adaLN=self.adaLN)
+
+  def _get_bias_dropout_scale(self):
+    if self.training:
+      return bias_dropout_add_scale_fused_train
+    else:
+      return  bias_dropout_add_scale_fused_inference
+
+  def forward(self, x, sigma, output_hidden_states=False):
+    all_hidden_states = []
+    x = self.vocab_embed(x)
+    if output_hidden_states:
+      all_hidden_states.append(x)
+    if self.causal:
+      t_cond = None
+    else:
+      t_cond = F.silu(self.sigma_map(sigma))
+
+    rotary_cos_sin = self.rotary_emb(x)
+    with torch.cuda.amp.autocast(
+      dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float32):
+      for i in range(len(self.blocks)):
+        x = self.blocks[i](x, rotary_cos_sin, c=t_cond)
+        if output_hidden_states:
+          all_hidden_states.append(x)
+      x = self.output_layer(x, c=t_cond)
+    return x, all_hidden_states
+
+
+class DUO(transformers.PreTrainedModel):
+  """HF-compatible model."""
+  config_class = DUOConfig
+  base_model_prefix = 'duo'
+
+  def __init__(self, config: DUOConfig):
+    super().__init__(config)
+    self.config = config
+    self.backbone = HFDIT(config)
+
+  def reset_kv_cache(self):
+    for block in self.backbone.blocks:
+      block.kv_cache = None
+
+  def forward(
+      self,
+      input_ids: torch.LongTensor = None,
+      timesteps: torch.FloatTensor = None,
+      output_hidden_states: typing.Optional[bool] = None,
+      return_dict: typing.Optional[bool] = None,
+  ) -> typing.Union[
+    torch.Tensor, typing.Tuple,
+    transformers.modeling_outputs.MaskedLMOutput]:
+    """HF-compatible forward method."""
+    output_hidden_states = (
+      output_hidden_states
+      if output_hidden_states is not None
+      else self.config.output_hidden_states
+    )
+    return_dict = return_dict \
+      if return_dict is not None \
+      else self.config.use_return_dict
+    
+    logits, all_hidden_states = self.backbone(
+      x=input_ids,
+      sigma=timesteps,
+      output_hidden_states=output_hidden_states,
+    )
+    if return_dict:
+      return transformers.modeling_outputs.MaskedLMOutput(
+        logits=logits,
+        hidden_states=all_hidden_states if output_hidden_states else None,
+        loss=None
+      )
+    elif output_hidden_states:
+      return logits, all_hidden_states
+    else:
+      return logits

checkpoints/official/langflow/lm1b/README.md

@@ -0,0 +1,61 @@
+---
+datasets:
+- dvruette/lm1b
+papers:
+- arxiv: 2604.11748
+language:
+- en
+library_name: transformers
+license: apache-2.0
+metrics:
+- perplexity
+pipeline_tag: text-generation
+---
+
+# LangFlow
+
+LangFlow is a continuous diffusion language model that operates in embedding space. Unlike discrete diffusion models (MDLM, SEDD, DUO), LangFlow performs diffusion directly on continuous token embeddings, enabling smoother denoising dynamics.
+
+For more details, please see our paper: [LangFlow: Continuous Diffusion Rivals Discrete in Language Modeling](https://arxiv.org/abs/2604.11748).
+
+
+## Using LangFlow
+
+To use the pre-trained model for text generation, use the following snippet:
+
+```python
+from transformers import AutoModelForMaskedLM, AutoTokenizer
+
+tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')
+model = AutoModelForMaskedLM.from_pretrained('Continuous-Rivals-Discrete/langflow-lm1b', trust_remote_code=True)
+
+# Generate samples
+samples = model.generate_samples(num_samples=5, num_steps=128)
+texts = tokenizer.batch_decode(samples, skip_special_tokens=True)
+for text in texts:
+    print(text)
+```
+
+## Model Details
+
+- **Architecture**: DiT (Diffusion Transformer) backbone with adaptive layer normalization
+- **Context Length**: 128 tokens
+- **Parameters**: ~130M parameters (similar to GPT-2 small)
+- **Training**: 1M steps on LM1B corpus
+- **Tokenizer**: bert-base-uncased tokenizer (30,522 vocab size)
+
+## Citation
+
+```
+@article{chen2026langflow,
+  title={LangFlow: Continuous Diffusion Rivals Discrete in Language Modeling},
+  author={Chen, Yuxin and Liang, Chumeng and Sui, Hangke and Guo, Ruihan and Cheng, Chaoran and You, Jiaxuan and Liu, Ge},
+  journal={arXiv preprint arXiv:2604.11748},
+  year={2026}
+}
+```
+
+## Model Card Contact
+
+Chumeng Liang (chumengl@illinois.edu)
+

checkpoints/official/langflow/lm1b/config.json

@@ -0,0 +1,29 @@
+{
+  "_name_or_path": "nealchen/langflow-lm1b",
+  "architectures": [
+    "LangFlow"
+  ],
+  "auto_map": {
+    "AutoConfig": "config.LangFlowConfig",
+    "AutoModelForMaskedLM": "model.LangFlow"
+  },
+  "model_type": "LangFlow",
+  "vocab_size": 30522,
+  "hidden_size": 768,
+  "cond_dim": 128,
+  "n_blocks": 12,
+  "n_heads": 12,
+  "dropout": 0.1,
+  "model_length": 128,
+  "use_normalized_embedding": true,
+  "embedding_norm_method": "layernorm",
+  "self_conditioning": true,
+  "use_bias": true,
+  "gumbel_loc": 4.723,
+  "gumbel_scale": 0.852,
+  "gumbel_cutoff": 1e-5,
+  "gumbel_entropy": 7.02,
+  "return_dict": true,
+  "torch_dtype": "float32",
+  "transformers_version": "4.38.2"
+}

checkpoints/official/langflow/lm1b/model.py

@@ -0,0 +1,560 @@
+"""HuggingFace model implementation for LangFlow.
+
+LangFlow is a continuous diffusion language model that operates in embedding space.
+"""
+
+import math
+import typing
+
+import einops
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import transformers
+
+from .config import LangFlowConfig
+
+
+# Flags required to enable jit fusion kernels
+torch._C._jit_set_profiling_mode(False)
+torch._C._jit_set_profiling_executor(False)
+torch._C._jit_override_can_fuse_on_cpu(True)
+torch._C._jit_override_can_fuse_on_gpu(True)
+
+
+def bias_dropout_add_scale(
+    x: torch.Tensor,
+    bias: typing.Optional[torch.Tensor],
+    scale: torch.Tensor,
+    residual: typing.Optional[torch.Tensor],
+    prob: float,
+    training: bool) -> torch.Tensor:
+    if bias is not None:
+        out = scale * F.dropout(x + bias, p=prob, training=training)
+    else:
+        out = scale * F.dropout(x, p=prob, training=training)
+
+    if residual is not None:
+        out = residual + out
+    return out
+
+
+@torch.jit.script
+def bias_dropout_add_scale_fused_train(
+    x: torch.Tensor,
+    bias: typing.Optional[torch.Tensor],
+    scale: torch.Tensor,
+    residual: typing.Optional[torch.Tensor],
+    prob: float) -> torch.Tensor:
+    return bias_dropout_add_scale(x, bias, scale, residual, prob, True)
+
+
+@torch.jit.script
+def bias_dropout_add_scale_fused_inference(
+    x: torch.Tensor,
+    bias: typing.Optional[torch.Tensor],
+    scale: torch.Tensor,
+    residual: typing.Optional[torch.Tensor],
+    prob: float) -> torch.Tensor:
+    return bias_dropout_add_scale(x, bias, scale, residual, prob, False)
+
+
+@torch.jit.script
+def modulate_fused(x: torch.Tensor,
+                   shift: torch.Tensor,
+                   scale: torch.Tensor) -> torch.Tensor:
+    return x * (1 + scale) + shift
+
+
+class Rotary(nn.Module):
+    def __init__(self, dim, base=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, seq_dim=1):
+        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)
+            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)
+            self.cos_cached[:, :, 2, :, :].fill_(1.)
+            self.sin_cached[:, :, 2, :, :].fill_(0.)
+        return self.cos_cached, self.sin_cached
+
+
+def _apply_rotary_emb(x, cos, sin):
+    # x: [batch, seqlen, nheads, headdim]
+    # cos, sin: [seqlen, headdim//2]
+    ro_dim = cos.shape[-1] * 2
+    # Expand to [1, seqlen, 1, ro_dim] for broadcasting
+    cos = torch.cat([cos, cos], dim=-1)[None, :, None, :]
+    sin = torch.cat([sin, sin], dim=-1)[None, :, None, :]
+    x_rot = x[..., :ro_dim]
+    x1, x2 = x_rot.chunk(2, dim=-1)
+    x_rotated = torch.cat([-x2, x1], dim=-1)
+    return torch.cat([x_rot * cos + x_rotated * sin, x[..., ro_dim:]], dim=-1)
+
+
+def split_and_apply_rotary_pos_emb(qkv, rotary_cos_sin):
+    with torch.autocast(device_type='cuda', enabled=False):
+        cos, sin = rotary_cos_sin
+        cos = cos.to(qkv.dtype)
+        sin = sin.to(qkv.dtype)
+        cos = cos[0, :, 0, 0, :cos.shape[-1]//2]
+        sin = sin[0, :, 0, 0, :sin.shape[-1]//2]
+        q, k, v = qkv.chunk(3, dim=2)
+        q = _apply_rotary_emb(q.squeeze(dim=2), cos, sin)
+        k = _apply_rotary_emb(k.squeeze(dim=2), cos, sin)
+        v = v.squeeze(dim=2)
+    return q, k, v
+
+
+def regular_attention_multi_headed(q, k, v):
+    attention_output = F.scaled_dot_product_attention(
+        query=q.transpose(1, 2),
+        key=k.transpose(1, 2),
+        value=v.transpose(1, 2),
+        attn_mask=None,
+        dropout_p=0.0,
+        is_causal=False)
+    attention_output = attention_output.transpose(1, 2)
+    return einops.rearrange(attention_output, 'b s h d -> b s (h d)')
+
+
+class LayerNorm(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones([dim]))
+        self.dim = dim
+
+    def forward(self, x):
+        with torch.autocast(device_type='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, frequency_embedding_size=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(t, dim, max_period=10000):
+        half = dim // 2
+        freqs = torch.exp(
+            -math.log(max_period)
+            * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device)
+            / half)
+        args = t[:, 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, t):
+        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+        t_emb = self.mlp(t_freq)
+        return t_emb
+
+
+class DDiTBlock(nn.Module):
+    def __init__(self, dim, n_heads, cond_dim, mlp_ratio=4, dropout=0.1):
+        super().__init__()
+        self.n_heads = n_heads
+
+        self.norm1 = LayerNorm(dim)
+        self.attn_qkv = nn.Linear(dim, 3 * dim, bias=False)
+        self.attn_out = nn.Linear(dim, dim, bias=False)
+
+        self.norm2 = LayerNorm(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.dropout = dropout
+
+        self.adaLN_modulation = nn.Linear(cond_dim, 6 * dim)
+        self.adaLN_modulation.weight.data.zero_()
+        self.adaLN_modulation.bias.data.zero_()
+
+    def _get_bias_dropout_scale(self):
+        if self.training:
+            return bias_dropout_add_scale_fused_train
+        else:
+            return bias_dropout_add_scale_fused_inference
+
+    def forward(self, x, rotary_cos_sin, c):
+        bias_dropout_scale_fn = self._get_bias_dropout_scale()
+
+        x_skip = x
+        x = self.norm1(x)
+
+        (shift_msa, scale_msa, gate_msa, shift_mlp,
+         scale_mlp, gate_mlp) = self.adaLN_modulation(c)[:, None].chunk(6, dim=2)
+        x = modulate_fused(x, shift_msa, scale_msa)
+
+        qkv = einops.rearrange(
+            self.attn_qkv(x),
+            'b s (three h d) -> b s three h d',
+            three=3,
+            h=self.n_heads)
+        q, k, v = split_and_apply_rotary_pos_emb(qkv, rotary_cos_sin)
+        x = regular_attention_multi_headed(q, k, v)
+
+        x = bias_dropout_scale_fn(self.attn_out(x), None, gate_msa, x_skip, self.dropout)
+        x = bias_dropout_scale_fn(
+            self.mlp(modulate_fused(self.norm2(x), shift_mlp, scale_mlp)),
+            None, gate_mlp, x, self.dropout)
+        return x
+
+
+def _normalize_embedding_layernorm(weight: torch.Tensor) -> torch.Tensor:
+    """Normalize embedding weights to unit norm per row, then scale by sqrt(dim)."""
+    normalized = F.normalize(weight.float(), dim=-1)
+    return (normalized * math.sqrt(weight.shape[-1])).to(weight.dtype)
+
+
+class EmbeddingLayer(nn.Module):
+    """Embedding layer with optional layernorm normalization."""
+    
+    def __init__(self, dim, vocab_dim, use_normalized_embedding=True):
+        super().__init__()
+        self.dim = dim
+        self.vocab_dim = vocab_dim
+        self.use_normalized_embedding = use_normalized_embedding
+        self.embedding = nn.Parameter(torch.empty((vocab_dim, dim)))
+        nn.init.kaiming_uniform_(self.embedding, a=math.sqrt(5))
+
+    def _get_embedding(self):
+        if self.use_normalized_embedding:
+            return _normalize_embedding_layernorm(self.embedding)
+        return self.embedding
+
+    def forward(self, x):
+        embedding = self._get_embedding()
+        if x.ndim == 2:
+            return embedding[x]
+        assert x.ndim == 3  # probabilities
+        return torch.einsum("blv,ve->ble", x.float(), embedding.float()).to(x.dtype)
+
+
+class DDiTFinalLayer(nn.Module):
+    def __init__(self, hidden_size, out_channels, cond_dim):
+        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, c):
+        x = self.norm_final(x)
+        shift, scale = self.adaLN_modulation(c)[:, None].chunk(2, dim=2)
+        x = modulate_fused(x, shift, scale)
+        x = self.linear(x)
+        return x
+
+
+class GumbelProposal(nn.Module):
+    """Learnable Gumbel distribution proposal for sampling gamma (log-SNR)."""
+    
+    def __init__(self, loc: float = 4.723, scale: float = 0.852, 
+                 cutoff: float = 1e-5, entropy: float = 7.02):
+        super().__init__()
+        self.loc = nn.Parameter(torch.tensor(loc))
+        self.scale = nn.Parameter(torch.tensor(scale))
+        self.cutoff = cutoff
+        self.entropy = nn.Parameter(torch.tensor(entropy))
+
+    def _get_distribution(self) -> torch.distributions.Gumbel:
+        return torch.distributions.Gumbel(self.loc, self.scale)
+
+    @property
+    def gamma_min(self) -> float:
+        return float(self.loc - math.log(-math.log(self.cutoff)) * self.scale)
+
+    @property
+    def gamma_max(self) -> float:
+        return float(self.loc - math.log(self.cutoff) * self.scale)
+
+    def forward(self, q: torch.Tensor) -> torch.Tensor:
+        """Convert uniform samples to gamma values via inverse CDF."""
+        gamma = self._get_distribution().icdf(q)
+        return gamma.clamp(min=self.gamma_min, max=self.gamma_max)
+
+    def log_pdf(self, gamma: torch.Tensor) -> torch.Tensor:
+        """Compute log probability density at gamma."""
+        return self._get_distribution().log_prob(gamma)
+
+
+class LangFlowBackbone(nn.Module):
+    """DiT backbone for LangFlow."""
+    
+    def __init__(self, config: LangFlowConfig):
+        super().__init__()
+        self.config = config
+        dim = config.hidden_size
+        cond_dim = config.cond_dim
+
+        self.vocab_embed = EmbeddingLayer(
+            dim, config.vocab_size,
+            use_normalized_embedding=config.use_normalized_embedding)
+        self.sigma_map = TimestepEmbedder(cond_dim)
+        self.rotary_emb = Rotary(dim // config.n_heads)
+
+        self.blocks = nn.ModuleList([
+            DDiTBlock(dim=dim, n_heads=config.n_heads, cond_dim=cond_dim, dropout=config.dropout)
+            for _ in range(config.n_blocks)
+        ])
+
+        self.output_layer = DDiTFinalLayer(
+            hidden_size=dim, out_channels=config.vocab_size, cond_dim=cond_dim)
+
+        # Self-conditioning projection
+        if config.self_conditioning:
+            self.self_cond_proj = nn.Linear(dim * 2, dim, bias=False)
+            nn.init.zeros_(self.self_cond_proj.weight)
+
+    def forward(self, x_embed, sigma, x_self_cond=None, output_hidden_states=False):
+        """Forward pass from embeddings.
+        
+        Args:
+            x_embed: [B, L, D] - Input embeddings (possibly noisy)
+            sigma: [B] - Gamma values (log-SNR)
+            x_self_cond: [B, L, D] - Self-conditioning embeddings (optional)
+            output_hidden_states: Whether to return all hidden states
+        
+        Returns:
+            logits: [B, L, vocab_size]
+            hidden_states: List of hidden states if output_hidden_states=True
+        """
+        all_hidden_states = []
+        x = x_embed
+        
+        if output_hidden_states:
+            all_hidden_states.append(x)
+
+        # Self-conditioning
+        if self.config.self_conditioning:
+            if x_self_cond is None:
+                x_self_cond = torch.zeros_like(x)
+            x = x + self.self_cond_proj(torch.cat([x, x_self_cond], dim=-1))
+
+        t_cond = F.silu(self.sigma_map(sigma))
+        rotary_cos_sin = self.rotary_emb(x)
+
+        with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
+            for block in self.blocks:
+                x = block(x, rotary_cos_sin, c=t_cond)
+                if output_hidden_states:
+                    all_hidden_states.append(x)
+            x = self.output_layer(x, c=t_cond)
+
+        return x, all_hidden_states
+
+
+class LangFlow(transformers.PreTrainedModel):
+    """HuggingFace-compatible LangFlow model.
+    
+    LangFlow is a continuous diffusion language model that operates in embedding space.
+    It uses a DiT (Diffusion Transformer) backbone with:
+    - Self-conditioning: uses previous predictions as additional input
+    - Bias (preconditioning): skip connection for improved generation
+    - Normalized embeddings: layernorm on embedding vectors
+    - Learnable Gumbel proposal for gamma (log-SNR) sampling
+    """
+    config_class = LangFlowConfig
+    base_model_prefix = "langflow"
+
+    def __init__(self, config: LangFlowConfig):
+        super().__init__(config)
+        self.config = config
+        self.backbone = LangFlowBackbone(config)
+        self.proposal = GumbelProposal(
+            loc=config.gumbel_loc,
+            scale=config.gumbel_scale,
+            cutoff=config.gumbel_cutoff,
+            entropy=config.gumbel_entropy)
+
+    def _get_embedding_matrix(self) -> torch.Tensor:
+        """Get the embedding matrix for bias skip connection."""
+        return self.backbone.vocab_embed._get_embedding()
+
+    def _embed_tokens(self, x: torch.Tensor) -> torch.Tensor:
+        """Embed tokens or probabilities to continuous embeddings."""
+        return self.backbone.vocab_embed(x)
+
+    def _forward_diffusion(self, x_embed: torch.Tensor, 
+                           gamma: torch.Tensor) -> torch.Tensor:
+        """Add noise to embeddings (forward diffusion process)."""
+        gamma = gamma.float()
+        alpha = torch.sigmoid(-gamma).sqrt()[:, None, None]
+        sigma = torch.sigmoid(gamma).sqrt()[:, None, None]
+        noise = torch.randn_like(x_embed)
+        return (x_embed * alpha + noise * sigma).to(x_embed.dtype)
+
+    def _euler_edm_step(self, z: torch.Tensor, x_pred: torch.Tensor,
+                        t: torch.Tensor, s: torch.Tensor) -> torch.Tensor:
+        """Single Euler step for EDM sampling."""
+        t_ = t.double()
+        s_ = s.double()
+        cur = z.double() * ((F.softplus(t_) - F.softplus(s_)) / 2).exp()
+        end = torch.sigmoid(-s_).sqrt() * x_pred.double()
+        z = end.lerp(cur, ((s_ - t_) / 2).exp()).to(z.dtype)
+        return z
+
+    def forward(
+        self,
+        input_ids: typing.Optional[torch.LongTensor] = None,
+        noisy_embeds: typing.Optional[torch.FloatTensor] = None,
+        timesteps: typing.Optional[torch.FloatTensor] = None,
+        x_self_cond: typing.Optional[torch.FloatTensor] = None,
+        output_hidden_states: typing.Optional[bool] = None,
+        return_dict: typing.Optional[bool] = None,
+    ) -> typing.Union[torch.Tensor, typing.Tuple, transformers.modeling_outputs.MaskedLMOutput]:
+        """Forward pass for LangFlow.
+        
+        Args:
+            input_ids: [B, L] - Token IDs (will be embedded and noised if timesteps provided)
+            noisy_embeds: [B, L, D] - Pre-noised embeddings (alternative to input_ids)
+            timesteps: [B] - Gamma values (log-SNR) for conditioning
+            x_self_cond: [B, L, D] - Self-conditioning embeddings
+            output_hidden_states: Whether to return hidden states
+            return_dict: Whether to return MaskedLMOutput
+        
+        Returns:
+            logits or MaskedLMOutput
+        """
+        output_hidden_states = output_hidden_states if output_hidden_states is not None else False
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # Get embeddings
+        if noisy_embeds is not None:
+            z = noisy_embeds
+        elif input_ids is not None:
+            x_embed = self._embed_tokens(input_ids)
+            if timesteps is not None:
+                z = self._forward_diffusion(x_embed, timesteps)
+            else:
+                z = x_embed
+        else:
+            raise ValueError("Either input_ids or noisy_embeds must be provided")
+
+        if timesteps is None:
+            # Use minimum gamma for clean input
+            timesteps = torch.full((z.shape[0],), self.proposal.gamma_min, device=z.device)
+
+        # Process sigma
+        sigma = timesteps
+        if sigma.ndim == 2:
+            sigma = sigma.mean(-1)
+
+        # Get model output
+        logits, all_hidden_states = self.backbone(
+            z, sigma, x_self_cond=x_self_cond, output_hidden_states=output_hidden_states)
+
+        # Add bias (preconditioning) skip connection
+        if self.config.use_bias:
+            c_skip = ((F.softplus(-sigma) - sigma) / 2).exp()
+            embedding = self._get_embedding_matrix()
+            skip_logits = torch.matmul(z.float(), embedding.t().float())
+            logits = logits + c_skip[:, None, None] * skip_logits.to(logits.dtype)
+
+        if return_dict:
+            return transformers.modeling_outputs.MaskedLMOutput(
+                logits=logits,
+                hidden_states=all_hidden_states if output_hidden_states else None,
+                loss=None)
+        elif output_hidden_states:
+            return logits, all_hidden_states
+        else:
+            return logits
+
+    @torch.no_grad()
+    def generate_samples(
+        self,
+        num_samples: int = 1,
+        seq_length: typing.Optional[int] = None,
+        num_steps: int = 128,
+        device: typing.Optional[torch.device] = None,
+    ) -> torch.LongTensor:
+        """Generate samples using Euler-EDM solver.
+        
+        Args:
+            num_samples: Number of samples to generate
+            seq_length: Sequence length (defaults to config.model_length)
+            num_steps: Number of denoising steps
+            device: Device to generate on
+        
+        Returns:
+            samples: [num_samples, seq_length] - Generated token IDs
+        """
+        if seq_length is None:
+            seq_length = self.config.model_length
+        if device is None:
+            device = next(self.parameters()).device
+
+        embed_dim = self.config.hidden_size
+        eps = 1e-5
+
+        # Initialize with Gaussian noise
+        z = torch.randn(num_samples, seq_length, embed_dim, device=device)
+
+        # Create gamma schedule from t=1-eps to t=eps
+        t = torch.linspace(1.0 - eps, eps, num_steps, device=device)
+        gamma = self.proposal(t)
+
+        # Self-conditioning state
+        x_self_cond = None
+
+        # Euler-EDM sampling loop
+        for i in range(len(gamma) - 1):
+            gamma_t = gamma[i]
+            gamma_s = gamma[i + 1]
+
+            # Get model prediction
+            gamma_expanded = gamma_t.unsqueeze(0).expand(num_samples)
+            logits = self.forward(
+                noisy_embeds=z,
+                timesteps=gamma_expanded,
+                x_self_cond=x_self_cond,
+                return_dict=False)
+
+            # Convert logits to embedding prediction
+            probs = F.softmax(logits.float(), dim=-1)
+            x_pred = self._embed_tokens(probs)
+
+            # Update self-conditioning
+            if self.config.self_conditioning:
+                x_self_cond = x_pred
+
+            # Euler step
+            z = self._euler_edm_step(z, x_pred, gamma_t, gamma_s)
+
+        # Final step: get logits and take argmax
+        gamma_final = gamma[-1]
+        gamma_expanded = gamma_final.unsqueeze(0).expand(num_samples)
+        logits = self.forward(
+            noisy_embeds=z,
+            timesteps=gamma_expanded,
+            x_self_cond=x_self_cond,
+            return_dict=False)
+        samples = logits.argmax(dim=-1)
+
+        return samples

checkpoints/official/langflow/owt/README.md

@@ -0,0 +1,61 @@
+---
+datasets:
+- Skylion007/openwebtext
+papers:
+- arxiv: 2604.11748
+language:
+- en
+library_name: transformers
+license: apache-2.0
+metrics:
+- perplexity
+pipeline_tag: text-generation
+---
+
+# LangFlow
+
+LangFlow is a continuous diffusion language model that operates in embedding space. Unlike discrete diffusion models (MDLM, SEDD, DUO), LangFlow performs diffusion directly on continuous token embeddings, enabling smoother denoising dynamics.
+
+For more details, please see our paper: [LangFlow: Continuous Diffusion Rivals Discrete in Language Modeling](https://arxiv.org/abs/2604.11748).
+
+
+## Using LangFlow
+
+To use the pre-trained model for text generation, use the following snippet:
+
+```python
+from transformers import AutoModelForMaskedLM, AutoTokenizer
+
+tokenizer = AutoTokenizer.from_pretrained('gpt2')
+model = AutoModelForMaskedLM.from_pretrained('chumengl/langflow-owt', trust_remote_code=True)
+
+# Generate samples
+samples = model.generate_samples(num_samples=5, num_steps=128)
+texts = tokenizer.batch_decode(samples, skip_special_tokens=True)
+for text in texts:
+    print(text)
+```
+
+## Model Details
+
+- **Architecture**: DiT (Diffusion Transformer) backbone with adaptive layer normalization
+- **Context Length**: 1024 tokens
+- **Parameters**: ~130M non-embedding parameters (similar to GPT-2 medium)
+- **Training**: 1M steps on OpenWebText corpus
+- **Tokenizer**: GPT-2 tokenizer (50,257 vocab size)
+
+## Citation
+
+```
+@article{chen2026langflow,
+  title={LangFlow: Continuous Diffusion Rivals Discrete in Language Modeling},
+  author={Chen, Yuxin and Liang, Chumeng and Sui, Hangke and Guo, Ruihan and Cheng, Chaoran and You, Jiaxuan and Liu, Ge},
+  journal={arXiv preprint arXiv:2604.11748},
+  year={2026}
+}
+```
+
+## Model Card Contact
+
+Chumeng Liang (chumengl@illinois.edu)
+

checkpoints/official/langflow/owt/__init__.py

@@ -0,0 +1,6 @@
+"""HuggingFace release package for LangFlow."""
+
+from .config import LangFlowConfig
+from .model import LangFlow
+
+__all__ = ["LangFlowConfig", "LangFlow"]

checkpoints/official/langflow/owt/model.py

@@ -0,0 +1,560 @@
+"""HuggingFace model implementation for LangFlow.
+
+LangFlow is a continuous diffusion language model that operates in embedding space.
+"""
+
+import math
+import typing
+
+import einops
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import transformers
+
+from .config import LangFlowConfig
+
+
+# Flags required to enable jit fusion kernels
+torch._C._jit_set_profiling_mode(False)
+torch._C._jit_set_profiling_executor(False)
+torch._C._jit_override_can_fuse_on_cpu(True)
+torch._C._jit_override_can_fuse_on_gpu(True)
+
+
+def bias_dropout_add_scale(
+    x: torch.Tensor,
+    bias: typing.Optional[torch.Tensor],
+    scale: torch.Tensor,
+    residual: typing.Optional[torch.Tensor],
+    prob: float,
+    training: bool) -> torch.Tensor:
+    if bias is not None:
+        out = scale * F.dropout(x + bias, p=prob, training=training)
+    else:
+        out = scale * F.dropout(x, p=prob, training=training)
+
+    if residual is not None:
+        out = residual + out
+    return out
+
+
+@torch.jit.script
+def bias_dropout_add_scale_fused_train(
+    x: torch.Tensor,
+    bias: typing.Optional[torch.Tensor],
+    scale: torch.Tensor,
+    residual: typing.Optional[torch.Tensor],
+    prob: float) -> torch.Tensor:
+    return bias_dropout_add_scale(x, bias, scale, residual, prob, True)
+
+
+@torch.jit.script
+def bias_dropout_add_scale_fused_inference(
+    x: torch.Tensor,
+    bias: typing.Optional[torch.Tensor],
+    scale: torch.Tensor,
+    residual: typing.Optional[torch.Tensor],
+    prob: float) -> torch.Tensor:
+    return bias_dropout_add_scale(x, bias, scale, residual, prob, False)
+
+
+@torch.jit.script
+def modulate_fused(x: torch.Tensor,
+                   shift: torch.Tensor,
+                   scale: torch.Tensor) -> torch.Tensor:
+    return x * (1 + scale) + shift
+
+
+class Rotary(nn.Module):
+    def __init__(self, dim, base=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, seq_dim=1):
+        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)
+            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)
+            self.cos_cached[:, :, 2, :, :].fill_(1.)
+            self.sin_cached[:, :, 2, :, :].fill_(0.)
+        return self.cos_cached, self.sin_cached
+
+
+def _apply_rotary_emb(x, cos, sin):
+    # x: [batch, seqlen, nheads, headdim]
+    # cos, sin: [seqlen, headdim//2]
+    ro_dim = cos.shape[-1] * 2
+    # Expand to [1, seqlen, 1, ro_dim] for broadcasting
+    cos = torch.cat([cos, cos], dim=-1)[None, :, None, :]
+    sin = torch.cat([sin, sin], dim=-1)[None, :, None, :]
+    x_rot = x[..., :ro_dim]
+    x1, x2 = x_rot.chunk(2, dim=-1)
+    x_rotated = torch.cat([-x2, x1], dim=-1)
+    return torch.cat([x_rot * cos + x_rotated * sin, x[..., ro_dim:]], dim=-1)
+
+
+def split_and_apply_rotary_pos_emb(qkv, rotary_cos_sin):
+    with torch.autocast(device_type='cuda', enabled=False):
+        cos, sin = rotary_cos_sin
+        cos = cos.to(qkv.dtype)
+        sin = sin.to(qkv.dtype)
+        cos = cos[0, :, 0, 0, :cos.shape[-1]//2]
+        sin = sin[0, :, 0, 0, :sin.shape[-1]//2]
+        q, k, v = qkv.chunk(3, dim=2)
+        q = _apply_rotary_emb(q.squeeze(dim=2), cos, sin)
+        k = _apply_rotary_emb(k.squeeze(dim=2), cos, sin)
+        v = v.squeeze(dim=2)
+    return q, k, v
+
+
+def regular_attention_multi_headed(q, k, v):
+    attention_output = F.scaled_dot_product_attention(
+        query=q.transpose(1, 2),
+        key=k.transpose(1, 2),
+        value=v.transpose(1, 2),
+        attn_mask=None,
+        dropout_p=0.0,
+        is_causal=False)
+    attention_output = attention_output.transpose(1, 2)
+    return einops.rearrange(attention_output, 'b s h d -> b s (h d)')
+
+
+class LayerNorm(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones([dim]))
+        self.dim = dim
+
+    def forward(self, x):
+        with torch.autocast(device_type='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, frequency_embedding_size=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(t, dim, max_period=10000):
+        half = dim // 2
+        freqs = torch.exp(
+            -math.log(max_period)
+            * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device)
+            / half)
+        args = t[:, 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, t):
+        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+        t_emb = self.mlp(t_freq)
+        return t_emb
+
+
+class DDiTBlock(nn.Module):
+    def __init__(self, dim, n_heads, cond_dim, mlp_ratio=4, dropout=0.1):
+        super().__init__()
+        self.n_heads = n_heads
+
+        self.norm1 = LayerNorm(dim)
+        self.attn_qkv = nn.Linear(dim, 3 * dim, bias=False)
+        self.attn_out = nn.Linear(dim, dim, bias=False)
+
+        self.norm2 = LayerNorm(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.dropout = dropout
+
+        self.adaLN_modulation = nn.Linear(cond_dim, 6 * dim)
+        self.adaLN_modulation.weight.data.zero_()
+        self.adaLN_modulation.bias.data.zero_()
+
+    def _get_bias_dropout_scale(self):
+        if self.training:
+            return bias_dropout_add_scale_fused_train
+        else:
+            return bias_dropout_add_scale_fused_inference
+
+    def forward(self, x, rotary_cos_sin, c):
+        bias_dropout_scale_fn = self._get_bias_dropout_scale()
+
+        x_skip = x
+        x = self.norm1(x)
+
+        (shift_msa, scale_msa, gate_msa, shift_mlp,
+         scale_mlp, gate_mlp) = self.adaLN_modulation(c)[:, None].chunk(6, dim=2)
+        x = modulate_fused(x, shift_msa, scale_msa)
+
+        qkv = einops.rearrange(
+            self.attn_qkv(x),
+            'b s (three h d) -> b s three h d',
+            three=3,
+            h=self.n_heads)
+        q, k, v = split_and_apply_rotary_pos_emb(qkv, rotary_cos_sin)
+        x = regular_attention_multi_headed(q, k, v)
+
+        x = bias_dropout_scale_fn(self.attn_out(x), None, gate_msa, x_skip, self.dropout)
+        x = bias_dropout_scale_fn(
+            self.mlp(modulate_fused(self.norm2(x), shift_mlp, scale_mlp)),
+            None, gate_mlp, x, self.dropout)
+        return x
+
+
+def _normalize_embedding_layernorm(weight: torch.Tensor) -> torch.Tensor:
+    """Normalize embedding weights to unit norm per row, then scale by sqrt(dim)."""
+    normalized = F.normalize(weight.float(), dim=-1)
+    return (normalized * math.sqrt(weight.shape[-1])).to(weight.dtype)
+
+
+class EmbeddingLayer(nn.Module):
+    """Embedding layer with optional layernorm normalization."""
+    
+    def __init__(self, dim, vocab_dim, use_normalized_embedding=True):
+        super().__init__()
+        self.dim = dim
+        self.vocab_dim = vocab_dim
+        self.use_normalized_embedding = use_normalized_embedding
+        self.embedding = nn.Parameter(torch.empty((vocab_dim, dim)))
+        nn.init.kaiming_uniform_(self.embedding, a=math.sqrt(5))
+
+    def _get_embedding(self):
+        if self.use_normalized_embedding:
+            return _normalize_embedding_layernorm(self.embedding)
+        return self.embedding
+
+    def forward(self, x):
+        embedding = self._get_embedding()
+        if x.ndim == 2:
+            return embedding[x]
+        assert x.ndim == 3  # probabilities
+        return torch.einsum("blv,ve->ble", x.float(), embedding.float()).to(x.dtype)
+
+
+class DDiTFinalLayer(nn.Module):
+    def __init__(self, hidden_size, out_channels, cond_dim):
+        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, c):
+        x = self.norm_final(x)
+        shift, scale = self.adaLN_modulation(c)[:, None].chunk(2, dim=2)
+        x = modulate_fused(x, shift, scale)
+        x = self.linear(x)
+        return x
+
+
+class GumbelProposal(nn.Module):
+    """Learnable Gumbel distribution proposal for sampling gamma (log-SNR)."""
+    
+    def __init__(self, loc: float = 4.723, scale: float = 0.852, 
+                 cutoff: float = 1e-5, entropy: float = 7.02):
+        super().__init__()
+        self.loc = nn.Parameter(torch.tensor(loc))
+        self.scale = nn.Parameter(torch.tensor(scale))
+        self.cutoff = cutoff
+        self.entropy = nn.Parameter(torch.tensor(entropy))
+
+    def _get_distribution(self) -> torch.distributions.Gumbel:
+        return torch.distributions.Gumbel(self.loc, self.scale)
+
+    @property
+    def gamma_min(self) -> float:
+        return float(self.loc - math.log(-math.log(self.cutoff)) * self.scale)
+
+    @property
+    def gamma_max(self) -> float:
+        return float(self.loc - math.log(self.cutoff) * self.scale)
+
+    def forward(self, q: torch.Tensor) -> torch.Tensor:
+        """Convert uniform samples to gamma values via inverse CDF."""
+        gamma = self._get_distribution().icdf(q)
+        return gamma.clamp(min=self.gamma_min, max=self.gamma_max)
+
+    def log_pdf(self, gamma: torch.Tensor) -> torch.Tensor:
+        """Compute log probability density at gamma."""
+        return self._get_distribution().log_prob(gamma)
+
+
+class LangFlowBackbone(nn.Module):
+    """DiT backbone for LangFlow."""
+    
+    def __init__(self, config: LangFlowConfig):
+        super().__init__()
+        self.config = config
+        dim = config.hidden_size
+        cond_dim = config.cond_dim
+
+        self.vocab_embed = EmbeddingLayer(
+            dim, config.vocab_size,
+            use_normalized_embedding=config.use_normalized_embedding)
+        self.sigma_map = TimestepEmbedder(cond_dim)
+        self.rotary_emb = Rotary(dim // config.n_heads)
+
+        self.blocks = nn.ModuleList([
+            DDiTBlock(dim=dim, n_heads=config.n_heads, cond_dim=cond_dim, dropout=config.dropout)
+            for _ in range(config.n_blocks)
+        ])
+
+        self.output_layer = DDiTFinalLayer(
+            hidden_size=dim, out_channels=config.vocab_size, cond_dim=cond_dim)
+
+        # Self-conditioning projection
+        if config.self_conditioning:
+            self.self_cond_proj = nn.Linear(dim * 2, dim, bias=False)
+            nn.init.zeros_(self.self_cond_proj.weight)
+
+    def forward(self, x_embed, sigma, x_self_cond=None, output_hidden_states=False):
+        """Forward pass from embeddings.
+        
+        Args:
+            x_embed: [B, L, D] - Input embeddings (possibly noisy)
+            sigma: [B] - Gamma values (log-SNR)
+            x_self_cond: [B, L, D] - Self-conditioning embeddings (optional)
+            output_hidden_states: Whether to return all hidden states
+        
+        Returns:
+            logits: [B, L, vocab_size]
+            hidden_states: List of hidden states if output_hidden_states=True
+        """
+        all_hidden_states = []
+        x = x_embed
+        
+        if output_hidden_states:
+            all_hidden_states.append(x)
+
+        # Self-conditioning
+        if self.config.self_conditioning:
+            if x_self_cond is None:
+                x_self_cond = torch.zeros_like(x)
+            x = x + self.self_cond_proj(torch.cat([x, x_self_cond], dim=-1))
+
+        t_cond = F.silu(self.sigma_map(sigma))
+        rotary_cos_sin = self.rotary_emb(x)
+
+        with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
+            for block in self.blocks:
+                x = block(x, rotary_cos_sin, c=t_cond)
+                if output_hidden_states:
+                    all_hidden_states.append(x)
+            x = self.output_layer(x, c=t_cond)
+
+        return x, all_hidden_states
+
+
+class LangFlow(transformers.PreTrainedModel):
+    """HuggingFace-compatible LangFlow model.
+    
+    LangFlow is a continuous diffusion language model that operates in embedding space.
+    It uses a DiT (Diffusion Transformer) backbone with:
+    - Self-conditioning: uses previous predictions as additional input
+    - Bias (preconditioning): skip connection for improved generation
+    - Normalized embeddings: layernorm on embedding vectors
+    - Learnable Gumbel proposal for gamma (log-SNR) sampling
+    """
+    config_class = LangFlowConfig
+    base_model_prefix = "langflow"
+
+    def __init__(self, config: LangFlowConfig):
+        super().__init__(config)
+        self.config = config
+        self.backbone = LangFlowBackbone(config)
+        self.proposal = GumbelProposal(
+            loc=config.gumbel_loc,
+            scale=config.gumbel_scale,
+            cutoff=config.gumbel_cutoff,
+            entropy=config.gumbel_entropy)
+
+    def _get_embedding_matrix(self) -> torch.Tensor:
+        """Get the embedding matrix for bias skip connection."""
+        return self.backbone.vocab_embed._get_embedding()
+
+    def _embed_tokens(self, x: torch.Tensor) -> torch.Tensor:
+        """Embed tokens or probabilities to continuous embeddings."""
+        return self.backbone.vocab_embed(x)
+
+    def _forward_diffusion(self, x_embed: torch.Tensor, 
+                           gamma: torch.Tensor) -> torch.Tensor:
+        """Add noise to embeddings (forward diffusion process)."""
+        gamma = gamma.float()
+        alpha = torch.sigmoid(-gamma).sqrt()[:, None, None]
+        sigma = torch.sigmoid(gamma).sqrt()[:, None, None]
+        noise = torch.randn_like(x_embed)
+        return (x_embed * alpha + noise * sigma).to(x_embed.dtype)
+
+    def _euler_edm_step(self, z: torch.Tensor, x_pred: torch.Tensor,
+                        t: torch.Tensor, s: torch.Tensor) -> torch.Tensor:
+        """Single Euler step for EDM sampling."""
+        t_ = t.double()
+        s_ = s.double()
+        cur = z.double() * ((F.softplus(t_) - F.softplus(s_)) / 2).exp()
+        end = torch.sigmoid(-s_).sqrt() * x_pred.double()
+        z = end.lerp(cur, ((s_ - t_) / 2).exp()).to(z.dtype)
+        return z
+
+    def forward(
+        self,
+        input_ids: typing.Optional[torch.LongTensor] = None,
+        noisy_embeds: typing.Optional[torch.FloatTensor] = None,
+        timesteps: typing.Optional[torch.FloatTensor] = None,
+        x_self_cond: typing.Optional[torch.FloatTensor] = None,
+        output_hidden_states: typing.Optional[bool] = None,
+        return_dict: typing.Optional[bool] = None,
+    ) -> typing.Union[torch.Tensor, typing.Tuple, transformers.modeling_outputs.MaskedLMOutput]:
+        """Forward pass for LangFlow.
+        
+        Args:
+            input_ids: [B, L] - Token IDs (will be embedded and noised if timesteps provided)
+            noisy_embeds: [B, L, D] - Pre-noised embeddings (alternative to input_ids)
+            timesteps: [B] - Gamma values (log-SNR) for conditioning
+            x_self_cond: [B, L, D] - Self-conditioning embeddings
+            output_hidden_states: Whether to return hidden states
+            return_dict: Whether to return MaskedLMOutput
+        
+        Returns:
+            logits or MaskedLMOutput
+        """
+        output_hidden_states = output_hidden_states if output_hidden_states is not None else False
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # Get embeddings
+        if noisy_embeds is not None:
+            z = noisy_embeds
+        elif input_ids is not None:
+            x_embed = self._embed_tokens(input_ids)
+            if timesteps is not None:
+                z = self._forward_diffusion(x_embed, timesteps)
+            else:
+                z = x_embed
+        else:
+            raise ValueError("Either input_ids or noisy_embeds must be provided")
+
+        if timesteps is None:
+            # Use minimum gamma for clean input
+            timesteps = torch.full((z.shape[0],), self.proposal.gamma_min, device=z.device)
+
+        # Process sigma
+        sigma = timesteps
+        if sigma.ndim == 2:
+            sigma = sigma.mean(-1)
+
+        # Get model output
+        logits, all_hidden_states = self.backbone(
+            z, sigma, x_self_cond=x_self_cond, output_hidden_states=output_hidden_states)
+
+        # Add bias (preconditioning) skip connection
+        if self.config.use_bias:
+            c_skip = ((F.softplus(-sigma) - sigma) / 2).exp()
+            embedding = self._get_embedding_matrix()
+            skip_logits = torch.matmul(z.float(), embedding.t().float())
+            logits = logits + c_skip[:, None, None] * skip_logits.to(logits.dtype)
+
+        if return_dict:
+            return transformers.modeling_outputs.MaskedLMOutput(
+                logits=logits,
+                hidden_states=all_hidden_states if output_hidden_states else None,
+                loss=None)
+        elif output_hidden_states:
+            return logits, all_hidden_states
+        else:
+            return logits
+
+    @torch.no_grad()
+    def generate_samples(
+        self,
+        num_samples: int = 1,
+        seq_length: typing.Optional[int] = None,
+        num_steps: int = 128,
+        device: typing.Optional[torch.device] = None,
+    ) -> torch.LongTensor:
+        """Generate samples using Euler-EDM solver.
+        
+        Args:
+            num_samples: Number of samples to generate
+            seq_length: Sequence length (defaults to config.model_length)
+            num_steps: Number of denoising steps
+            device: Device to generate on
+        
+        Returns:
+            samples: [num_samples, seq_length] - Generated token IDs
+        """
+        if seq_length is None:
+            seq_length = self.config.model_length
+        if device is None:
+            device = next(self.parameters()).device
+
+        embed_dim = self.config.hidden_size
+        eps = 1e-5
+
+        # Initialize with Gaussian noise
+        z = torch.randn(num_samples, seq_length, embed_dim, device=device)
+
+        # Create gamma schedule from t=1-eps to t=eps
+        t = torch.linspace(1.0 - eps, eps, num_steps, device=device)
+        gamma = self.proposal(t)
+
+        # Self-conditioning state
+        x_self_cond = None
+
+        # Euler-EDM sampling loop
+        for i in range(len(gamma) - 1):
+            gamma_t = gamma[i]
+            gamma_s = gamma[i + 1]
+
+            # Get model prediction
+            gamma_expanded = gamma_t.unsqueeze(0).expand(num_samples)
+            logits = self.forward(
+                noisy_embeds=z,
+                timesteps=gamma_expanded,
+                x_self_cond=x_self_cond,
+                return_dict=False)
+
+            # Convert logits to embedding prediction
+            probs = F.softmax(logits.float(), dim=-1)
+            x_pred = self._embed_tokens(probs)
+
+            # Update self-conditioning
+            if self.config.self_conditioning:
+                x_self_cond = x_pred
+
+            # Euler step
+            z = self._euler_edm_step(z, x_pred, gamma_t, gamma_s)
+
+        # Final step: get logits and take argmax
+        gamma_final = gamma[-1]
+        gamma_expanded = gamma_final.unsqueeze(0).expand(num_samples)
+        logits = self.forward(
+            noisy_embeds=z,
+            timesteps=gamma_expanded,
+            x_self_cond=x_self_cond,
+            return_dict=False)
+        samples = logits.argmax(dim=-1)
+
+        return samples

checkpoints/official/mdlm/owt/README.md

@@ -0,0 +1,67 @@
+---
+library_name: transformers
+license: apache-2.0
+language:
+- en
+datasets:
+- Skylion007/openwebtext
+metrics:
+- perplexity
+---
+
+## Using MDLM
+To use the pre-trained model for masked language modeling, use the following snippet:
+```python
+from transformers import AutoModelForMaskedLM, AutoTokenizer
+
+# See the `MDLM` collection page on the hub for list of available models.
+tokenizer = transformers.AutoTokenizer.from_pretrained('gpt2')
+model_name = 'kuleshov-group/mdlm-owt'
+model = AutoModelForMaskedLM.from_pretrained(model_name)
+```
+
+For more details, please see our github repository: [MDLM](https://github.com/kuleshov-group/mdlm)
+
+## Model Details
+The model, which has a context length of `1024` and is similar in size to GPT2-medium with approximately `130 million` non-embedding parameters,
+was trained using a forward diffusion process that generates inputs varying from fully masked to fully unmasked. Its objective is to
+reconstruct the original input from these varying levels of masking, outputting logits in the process.
+The training regimen comprised one million steps on the OpenWebText corpus, involving the processing of a total of `33 billion` tokens.
+
+For more details, please see our paper: [Simple and Effective Masked Diffusion Language Models](http://arxiv.org/abs/2406.07524).
+
+
+
+## Citation
+
+<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
+Please cite our work using the bibtex below:
+
+**BibTeX:**
+
+```
+@inproceedings{
+    sahoo2024simple,
+    title={Simple and Effective Masked Diffusion Language Models},
+    author={Subham Sekhar Sahoo and Marianne Arriola and Aaron Gokaslan and Edgar Mariano Marroquin and Alexander M Rush and Yair Schiff and Justin T Chiu and Volodymyr Kuleshov},
+    booktitle={The Thirty-eighth Annual Conference on Neural Information Processing Systems},
+    year={2024},
+    url={https://openreview.net/forum?id=L4uaAR4ArM}
+}
+```
+
+**APA:**
+
+```
+@software{Sahoo_Simple_and_Effective_2024,
+author = {Sahoo, Subham Sekhar and Arriola, Marianne and Schiff, Yair and Gokaslan, Aaron and Marroquin, Edgar and Chiu, Justin T and Rush, Alexander and Kuleshov, Volodymyr},
+doi = {10.48550/arXiv.2406.07524},
+month = jun,
+title = {{Simple and Effective Masked Diffusion Language Models}},
+version = {arXiv:2406.07524v1},
+year = {2024}
+}
+```
+
+## Model Card Contact
+Subham Sekhar Sahoo (subbham@gmail.com)

checkpoints/official/mdlm/owt/config.json

@@ -0,0 +1,22 @@
+{
+  "_name_or_path": "kuleshov-group/mdlm-owt",
+  "architectures": [
+    "MDLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_mdlm.MDLMConfig",
+    "AutoModelForMaskedLM": "modeling_mdlm.MDLM"
+  },
+  "cond_dim": 128,
+  "dropout": 0.1,
+  "hidden_dim": 768,
+  "model_length": 1024,
+  "model_type": "mdlm",
+  "n_blocks": 12,
+  "n_heads": 12,
+  "return_dict": false,
+  "time_conditioning": false,
+  "torch_dtype": "float32",
+  "transformers_version": "4.38.2",
+  "vocab_size": 50258
+}

checkpoints/official/mdlm/owt/configuration_mdlm.py

@@ -0,0 +1,31 @@
+"""MDLM config for Hugging Face.
+
+"""
+
+import transformers
+
+
+class MDLMConfig(transformers.PretrainedConfig):
+  """Hugging Face configuration class for MDLM."""
+  model_type = "mdlm"
+
+  def __init__(
+    self,
+    vocab_size: int = 50258,
+    model_length: int = 1024,
+    hidden_dim: int = 768,
+    cond_dim: int = 129,
+    n_blocks: int = 12,
+    n_heads: int = 12,
+    dropout: float = 0.1,
+    time_conditioning: bool = False,
+    ** kwargs):
+    super().__init__(**kwargs)
+    self.vocab_size = vocab_size
+    self.model_length = model_length
+    self.hidden_dim = hidden_dim
+    self.cond_dim = cond_dim
+    self.n_blocks = n_blocks
+    self.n_heads = n_heads
+    self.dropout = dropout
+    self.time_conditioning = time_conditioning

checkpoints/official/mdlm/owt/modeling_mdlm.py

@@ -0,0 +1,450 @@
+"""MDLM model for Hugging Face.
+
+"""
+import itertools
+import math
+import typing
+
+import flash_attn
+import flash_attn.layers.rotary
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import transformers
+from einops import rearrange
+from transformers import modeling_outputs
+
+from .configuration_mdlm import MDLMConfig
+
+# Flags required to enable jit fusion kernels
+torch._C._jit_set_profiling_mode(False)
+torch._C._jit_set_profiling_executor(False)
+torch._C._jit_override_can_fuse_on_cpu(True)
+torch._C._jit_override_can_fuse_on_gpu(True)
+
+
+def bias_dropout_add_scale(
+    x: torch.Tensor,
+    bias: typing.Optional[torch.Tensor],
+    scale: torch.Tensor,
+    residual: typing.Optional[torch.Tensor],
+    prob: float,
+    training: bool) -> torch.Tensor:
+  if bias is not None:
+    out = scale * F.dropout(x + bias, p=prob,
+                            training=training)
+  else:
+    out = scale * F.dropout(x, p=prob, training=training)
+
+  if residual is not None:
+    out = residual + out
+  return out
+
+
+def get_bias_dropout_add_scale(training):
+  def _bias_dropout_add(x, bias, scale, residual, prob):
+    return bias_dropout_add_scale(
+      x, bias, scale, residual, prob, training)
+
+  return _bias_dropout_add
+
+
+# function overload
+def modulate(x: torch.Tensor,
+             shift: torch.Tensor,
+             scale: torch.Tensor) -> torch.Tensor:
+  return x * (1 + scale) + shift
+
+
+@torch.jit.script
+def bias_dropout_add_scale_fused_train(
+    x: torch.Tensor,
+    bias: typing.Optional[torch.Tensor],
+    scale: torch.Tensor,
+    residual: typing.Optional[torch.Tensor],
+    prob: float) -> torch.Tensor:
+  return bias_dropout_add_scale(
+    x, bias, scale, residual, prob, True)
+
+
+@torch.jit.script
+def bias_dropout_add_scale_fused_inference(
+    x: torch.Tensor,
+    bias: typing.Optional[torch.Tensor],
+    scale: torch.Tensor,
+    residual: typing.Optional[torch.Tensor],
+    prob: float) -> torch.Tensor:
+  return bias_dropout_add_scale(
+    x, bias, scale, residual, prob, False)
+
+
+@torch.jit.script
+def modulate_fused(x: torch.Tensor,
+                   shift: torch.Tensor,
+                   scale: torch.Tensor) -> torch.Tensor:
+  return modulate(x, shift, scale)
+
+
+class Rotary(torch.nn.Module):
+  def __init__(self, dim, base=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, seq_dim=1):
+    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.)
+      self.sin_cached[:, :, 2, :, :].fill_(0.)
+
+    return self.cos_cached, self.sin_cached
+
+
+def rotate_half(x):
+  x1, x2 = x[..., : x.shape[-1] // 2], x[...,
+                                       x.shape[-1] // 2:]
+  return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(qkv, cos, sin):
+  cos = cos[0, :, 0, 0, :cos.shape[-1] // 2]
+  sin = sin[0, :, 0, 0, :sin.shape[-1] // 2]
+  return flash_attn.layers.rotary.apply_rotary_emb_qkv_(qkv,
+                                                        cos,
+                                                        sin)
+# function overload
+def modulate(x, shift, scale):
+  return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+############################################################
+#                          Layers                          #
+############################################################
+class LayerNorm(nn.Module):
+  def __init__(self, dim):
+    super().__init__()
+    self.weight = nn.Parameter(torch.ones([dim]))
+    self.dim = dim
+
+  def forward(self, x):
+    with torch.cuda.amp.autocast(enabled=False):
+      x = F.layer_norm(x.float(), [self.dim])
+    return x * self.weight[None, None, :]
+
+
+def residual_linear(x, W, x_skip, residual_scale):
+  """x_skip + residual_scale * W @ x"""
+  dim_out, dim_in = W.shape[0], W.shape[1]
+  return torch.addmm(
+    x_skip.view(-1, dim_out),
+    x.view(-1, dim_in),
+    W.T,
+    alpha=residual_scale).view(*x.shape[:-1], dim_out)
+
+
+############################################################
+#     Embedding Layers for Timesteps and Class Labels      #
+############################################################
+class TimestepEmbedder(nn.Module):
+  """
+  Embeds scalar timesteps into vector representations.
+  """
+
+  def __init__(self, hidden_size,
+               frequency_embedding_size=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(t, dim, max_period=10000):
+    """
+    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.
+    """
+    # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
+    half = dim // 2
+    freqs = torch.exp(
+      - math.log(max_period)
+      * torch.arange(start=0, end=half, dtype=torch.float32)
+      / half).to(device=t.device)
+    args = t[:, 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, t):
+    t_freq = self.timestep_embedding(t,
+                                     self.frequency_embedding_size)
+    t_emb = self.mlp(t_freq)
+    return t_emb
+
+
+class LabelEmbedder(nn.Module):
+  """Embeds class labels into vector representations.
+
+  Also handles label dropout for classifier-free guidance.
+  """
+
+  def __init__(self, num_classes, cond_size):
+    super().__init__()
+    self.embedding_table = nn.Embedding(num_classes + 1,
+                                        cond_size)
+    self.num_classes = num_classes
+
+    # TODO think of initializing with 0.02 std deviation like in original DiT paper
+
+  def forward(self, labels):
+    embeddings = self.embedding_table(labels)
+    return embeddings
+
+
+############################################################
+#                        Core Model                        #
+############################################################
+
+
+class DDiTBlock(nn.Module):
+  def __init__(self, dim, n_heads, cond_dim, mlp_ratio=4,
+               dropout=0.1):
+    super().__init__()
+    self.n_heads = n_heads
+
+    self.norm1 = LayerNorm(dim)
+    self.attn_qkv = nn.Linear(dim, 3 * dim, bias=False)
+    self.attn_out = nn.Linear(dim, dim, bias=False)
+    self.dropout1 = nn.Dropout(dropout)
+
+    self.norm2 = LayerNorm(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.dropout2 = nn.Dropout(dropout)
+    self.dropout = dropout
+
+    self.adaLN_modulation = nn.Linear(cond_dim, 6 * dim,
+                                      bias=True)
+    self.adaLN_modulation.weight.data.zero_()
+    self.adaLN_modulation.bias.data.zero_()
+
+  def _get_bias_dropout_scale(self):
+    if self.training:
+      return bias_dropout_add_scale_fused_train
+    else:
+      return bias_dropout_add_scale_fused_inference
+
+  def forward(self, x, rotary_cos_sin, c, seqlens=None):
+    batch_size, seq_len = x.shape[0], x.shape[1]
+
+    bias_dropout_scale_fn = self._get_bias_dropout_scale()
+
+    (shift_msa, scale_msa, gate_msa, shift_mlp,
+     scale_mlp, gate_mlp) = self.adaLN_modulation(c)[:,
+                            None].chunk(6, dim=2)
+
+    # attention operation
+    x_skip = x
+    x = modulate_fused(self.norm1(x), shift_msa, scale_msa)
+
+    qkv = self.attn_qkv(x)
+    qkv = rearrange(qkv,
+                    'b s (three h d) -> b s three h d',
+                    three=3,
+                    h=self.n_heads)
+    with torch.cuda.amp.autocast(enabled=False):
+      cos, sin = rotary_cos_sin
+      qkv = apply_rotary_pos_emb(
+        qkv, cos.to(qkv.dtype), sin.to(qkv.dtype))
+    qkv = rearrange(qkv, 'b s ... -> (b s) ...')
+    if seqlens is None:
+      cu_seqlens = torch.arange(
+        0, (batch_size + 1) * seq_len, step=seq_len,
+        dtype=torch.int32, device=qkv.device)
+    else:
+      cu_seqlens = seqlens.cumsum(-1)
+    x = flash_attn.flash_attn_interface.flash_attn_varlen_qkvpacked_func(
+      qkv, cu_seqlens, seq_len, 0., causal=False)
+
+    x = rearrange(x, '(b s) h d -> b s (h d)', b=batch_size)
+
+    x = bias_dropout_scale_fn(self.attn_out(x),
+                              None,
+                              gate_msa,
+                              x_skip,
+                              self.dropout)
+
+    # mlp operation
+    x = bias_dropout_scale_fn(
+      self.mlp(modulate_fused(
+        self.norm2(x), shift_mlp, scale_mlp)),
+      None, gate_mlp, x, self.dropout)
+    return x
+
+
+class EmbeddingLayer(nn.Module):
+  def __init__(self, dim, vocab_dim):
+    super().__init__()
+    self.embedding = nn.Parameter(
+      torch.empty((vocab_dim, dim)))
+    torch.nn.init.kaiming_uniform_(self.embedding,
+                                   a=math.sqrt(5))
+
+  def forward(self, x):
+    return self.embedding[x]
+
+
+class DDitFinalLayer(nn.Module):
+  def __init__(self, hidden_size, out_channels, cond_dim):
+    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, c):
+    shift, scale = self.adaLN_modulation(c)[:, None].chunk(
+      2, dim=2)
+    x = modulate_fused(self.norm_final(x), shift, scale)
+    x = self.linear(x)
+    return x
+
+
+class DITBackbone(nn.Module):
+  def __init__(
+      self,
+      config: MDLMConfig):
+    super().__init__()
+
+    self.config = config
+    self.vocab_size = config.vocab_size
+
+    self.vocab_embed = EmbeddingLayer(
+      config.hidden_dim,
+      config.vocab_size)
+    self.sigma_map = TimestepEmbedder(
+      config.cond_dim)
+    self.rotary_emb = Rotary(
+      config.hidden_dim // config.n_heads)
+
+    blocks = []
+    for _ in range(config.n_blocks):
+      blocks.append(DDiTBlock(config.hidden_dim,
+                              config.n_heads,
+                              config.cond_dim,
+                              dropout=config.dropout))
+    self.blocks = nn.ModuleList(blocks)
+
+    self.output_layer = DDitFinalLayer(
+      config.hidden_dim,
+      config.vocab_size,
+      config.cond_dim)
+
+  def _get_bias_dropout_scale(self):
+    if self.training:
+      return bias_dropout_add_scale_fused_train
+    else:
+      return bias_dropout_add_scale_fused_inference
+
+  def forward(self, indices, sigma,
+              output_hidden_states=False):
+    if not self.config.time_conditioning:
+      sigma = torch.zeros_like(sigma)
+    all_hidden_states = []
+    x = self.vocab_embed(indices)
+    if output_hidden_states:
+      all_hidden_states.append(x)
+    c = F.silu(self.sigma_map(sigma))
+
+    rotary_cos_sin = self.rotary_emb(x)
+
+    with torch.cuda.amp.autocast(dtype=torch.bfloat16):
+      for i in range(len(self.blocks)):
+        x = self.blocks[i](x, rotary_cos_sin, c,
+                           seqlens=None)
+        if output_hidden_states:
+          all_hidden_states.append(x)
+      logits = self.output_layer(x, c)
+    return logits, all_hidden_states
+
+class MDLM(transformers.PreTrainedModel):
+  """HF-compatible model."""
+  config_class = MDLMConfig
+  base_model_prefix = "mdlm"
+
+  def __init__(
+    self,
+    config: MDLMConfig):
+    super().__init__(config)
+    self.backbone = DITBackbone(config)
+
+  def forward(
+      self,
+      input_ids: torch.LongTensor = None,
+      timesteps: torch.FloatTensor = None,
+      output_hidden_states: typing.Optional[bool] = None,
+      return_dict: typing.Optional[bool] = None,
+  ) -> typing.Union[
+    torch.Tensor, typing.Tuple,
+    modeling_outputs.MaskedLMOutput]:
+    """HF-compatible forward method."""
+    output_hidden_states = (
+      output_hidden_states
+      if output_hidden_states is not None
+      else self.config.output_hidden_states
+    )
+    return_dict = return_dict \
+      if return_dict is not None \
+      else self.config.use_return_dict
+
+    logits, all_hidden_states = self.backbone(
+      indices=input_ids,
+      sigma=timesteps,
+      output_hidden_states=output_hidden_states
+    )
+    if return_dict:
+      return modeling_outputs.MaskedLMOutput(
+        logits=logits,
+        hidden_states=all_hidden_states if output_hidden_states else None,
+        loss=None
+      )
+    elif output_hidden_states:
+      return logits, all_hidden_states
+    else:
+      return logits

checkpoints/quarantine/20260802T040549Z/official/fmlm/lm1b/README.md

@@ -0,0 +1,45 @@
+---
+tags:
+- model_hub_mixin
+- pytorch_model_hub_mixin
+license: apache-2.0
+language:
+- en
+base_model:
+- david3684/FLM-B-LM1B
+---
+  
+  ## FMLM (LM1B)
+
+  For more information and implementation details, visit our github repository: [FLM](https://github.com/david3684/flm).
+
+  ## Model Details
+
+  The model has a context length of 128, distilled from 1M steps trained FLM checkpoint trained on the One Billion Words (LM1B) dataset.
+  
+  For more details, please see our paper: 
+  - [arXiv:2602.16813](https://arxiv.org/abs/2602.16813)
+  - [Huggingface Papers](https://huggingface.co/papers/2602.16813)
+
+  ## Citation
+
+  Please cite our work using the bibtex below:
+
+  **BibTeX:**
+
+
+  ```bibtex
+  @article{lee2026flow,
+      title={Flow Map Language Models: One-step Language Modeling via Continuous Denoising},
+      author={Chanhyuk Lee and Jaehoon Yoo and Manan Agarwal
+              and Sheel Shah and Jerry Huang
+              and Aditi Raghunathan and Seunghoon Hong
+              and Nicholas M. Boffi and Jinwoo Kim},
+      journal={arXiv preprint arXiv:2602.16813},
+      year={2026},
+  }
+  ```
+
+  ## Model Card Contact
+
+  Chanhyuk Lee (chanhyuk-lee@kaist.ac.kr)

checkpoints/quarantine/20260802T040549Z/official/fmlm/lm1b/config.json

@@ -0,0 +1,221 @@
+{
+  "config": {
+    "algo": {
+      "T": 0,
+      "add_boundary": "fixed",
+      "backbone": "dit",
+      "backprop_entmax_temp_lsd": false,
+      "boundary_prob": 32,
+      "causal_attention": false,
+      "diagonal_fraction": 0.5,
+      "distillation_method": "PSD",
+      "double_temb": true,
+      "entmax_temp_lsd": 2.0,
+      "initialize_student_from_teacher": false,
+      "learnable_loss_weighting": false,
+      "loss_scale": "none",
+      "loss_type": null,
+      "name": "fmlm",
+      "offdiagonal_sampling": "uniform_diff",
+      "parameterization": "mean",
+      "pred_log_interval": 2000,
+      "rescale_offdiag_loss_psd": false,
+      "subs_masking": false,
+      "t_max": 1.0,
+      "t_min": 0.0,
+      "teacher_path": "/home/david3684/ours_flow_ce_1m.ckpt",
+      "time_condition": "alpha_t",
+      "time_conditioning": true,
+      "use_discrete_schedule": true,
+      "use_mse_loss_psd": false,
+      "use_simulated_entropy": false,
+      "use_teacher_for_D_t_lsd": false
+    },
+    "callbacks": {
+      "checkpoint_every_n_steps": {
+        "_target_": "lightning.pytorch.callbacks.ModelCheckpoint",
+        "auto_insert_metric_name": false,
+        "dirpath": "/home/david3684/projects/flm/hf_upload/checkpoints",
+        "every_n_train_steps": 25000,
+        "save_last": false,
+        "save_top_k": -1,
+        "verbose": true
+      },
+      "checkpoint_monitor": {
+        "_target_": "lightning.pytorch.callbacks.ModelCheckpoint",
+        "auto_insert_metric_name": false,
+        "dirpath": "/home/david3684/projects/flm/hf_upload/checkpoints",
+        "filename": "best_nll",
+        "mode": "min",
+        "monitor": "val/nll",
+        "save_last": true,
+        "save_top_k": 1,
+        "verbose": true
+      },
+      "learning_rate_monitor": {
+        "_target_": "lightning.pytorch.callbacks.LearningRateMonitor",
+        "logging_interval": "step"
+      }
+    },
+    "checkpointing": {
+      "resume_ckpt_path": "/home/david3684/projects/flm/hf_upload/checkpoints/last.ckpt",
+      "resume_from_ckpt": false,
+      "save_dir": "/home/david3684/projects/flm/hf_upload"
+    },
+    "data": {
+      "cache_dir": "/home/david3684/datasets/lm1b",
+      "insert_train_eos": true,
+      "insert_valid_eos": true,
+      "streaming": false,
+      "tokenizer_name_or_path": "bert-base-uncased",
+      "train": "lm1b",
+      "valid": "lm1b",
+      "wrap": true
+    },
+    "eval": {
+      "checkpoint_path": "",
+      "compute_generative_perplexity": true,
+      "compute_perplexity_on_sanity": false,
+      "disable_ema": false,
+      "ema_decay": null,
+      "gen_ppl_eval_model_name_or_path": "gpt2-large",
+      "generate_samples": true,
+      "generated_samples_path": "/home/david3684/projects/flm/hf_upload/samples.json",
+      "perplexity_batch_size": 8
+    },
+    "is_di4c": false,
+    "is_di4c_deterministic": true,
+    "loader": {
+      "batch_size": 128,
+      "eval_batch_size": 128,
+      "eval_global_batch_size": 512,
+      "global_batch_size": 512,
+      "num_workers": 0,
+      "pin_memory": true
+    },
+    "lr_scheduler": {
+      "_target_": "transformers.get_constant_schedule_with_warmup",
+      "num_warmup_steps": 2500
+    },
+    "mode": "train",
+    "model": {
+      "cond_dim": 128,
+      "dropout": 0.1,
+      "hidden_size": 768,
+      "length": 128,
+      "n_blocks": 12,
+      "n_heads": 12,
+      "name": "small",
+      "scale_by_sigma": true,
+      "tie_word_embeddings": false,
+      "type": "ddit",
+      "vocab_lookup": true
+    },
+    "noise": {
+      "denoiser_latent_conditioning": -1,
+      "eps": 0,
+      "freeze_decoder": false,
+      "freeze_encoder": false,
+      "parameterization": "log-linear",
+      "type": "log-linear"
+    },
+    "optim": {
+      "beta1": 0.9,
+      "beta2": 0.999,
+      "eps": 1e-08,
+      "ln_tune": "none",
+      "lr": 0.0003,
+      "muon_lr": 0.03,
+      "muon_momentum": 0.95,
+      "muon_weight_decay": 0.01,
+      "use_muon": false,
+      "weight_decay": 0
+    },
+    "prior": {
+      "latent_height": 0,
+      "latent_width": 0,
+      "type": "none"
+    },
+    "sampling": {
+      "argmax_correction": false,
+      "autoguidance_weight": 0.0,
+      "hard_start": false,
+      "noise_removal": "uniform_alpha",
+      "num_reflow_samples": 50000,
+      "num_sample_batches": 1,
+      "num_sample_log": 2,
+      "num_strides": 1,
+      "p_nucleus": 1.0,
+      "predictor": "ancestral",
+      "semi_ar": false,
+      "solver": "euler",
+      "steps": [
+        1,
+        2,
+        4,
+        8,
+        16,
+        32,
+        64,
+        128
+      ],
+      "stride_length": 1,
+      "tau_log10": -1.0,
+      "temperature": 1.0,
+      "use_float64": true
+    },
+    "seed": 1,
+    "strategy": {
+      "_target_": "lightning.pytorch.strategies.DDPStrategy",
+      "find_unused_parameters": false
+    },
+    "trainer": {
+      "_target_": "lightning.Trainer",
+      "accelerator": "cuda",
+      "accumulate_grad_batches": 1,
+      "devices": 8,
+      "gradient_clip_val": 1.0,
+      "limit_train_batches": 1.0,
+      "limit_val_batches": 10,
+      "log_every_n_steps": 100,
+      "max_steps": 1000000,
+      "num_nodes": 1,
+      "num_sanity_val_steps": 2,
+      "precision": "bf16",
+      "val_check_interval": 10000
+    },
+    "training": {
+      "antithetic_sampling": true,
+      "change_of_variables": false,
+      "double_ema": 0.0,
+      "ema": "0.9999, 0.999, 0.99, 0.95",
+      "finetune_path": "",
+      "importance_sampling": false,
+      "load_ema": false,
+      "loss_precision": "bf16",
+      "loss_type": "sd",
+      "loss_weighting": false,
+      "not_sampling_t": false,
+      "pred_type": "x0",
+      "sampling_eps": 0.001,
+      "t_curriculum_max": 1.0,
+      "t_curriculum_min": 0.0,
+      "t_curriculum_steps": 100000,
+      "torch_compile_mode": "default",
+      "use_torch_compile": true
+    },
+    "wandb": {
+      "group": null,
+      "id": "27d989c8",
+      "job_type": null,
+      "name": "lm1b_PSD_teacher_distill_DF_0.5_BP_32_SE_False_IT_False_EDMW_False_27d989c8",
+      "notes": "DOS",
+      "project": "dos_distill",
+      "tags": [
+        "lm1b",
+        "dos_denoiser"
+      ]
+    }
+  },
+  "vocab_size": 30522
+}

checkpoints/quarantine/20260802T040557Z/official/flm/owt/README.md

@@ -0,0 +1,45 @@
+---
+tags:
+- model_hub_mixin
+- pytorch_model_hub_mixin
+license: apache-2.0
+datasets:
+- Skylion007/openwebtext
+language:
+- en
+---
+  
+  ## FLM (OpenWebText)
+
+  For more information and implementation details, visit our github repository: [FLM](https://github.com/david3684/flm).
+
+  ## Model Details
+
+  The model has a context length of 1024, trained for 1M steps on the OpenWebText dataset.
+
+  For more details, please see our paper: 
+  - [arXiv:2602.16813](https://arxiv.org/abs/2602.16813)
+  - [Huggingface Papers](https://huggingface.co/papers/2602.16813)
+
+  ## Citation
+
+  Please cite our work using the bibtex below:
+
+  **BibTeX:**
+
+
+  ```bibtex
+  @article{lee2026flow,
+      title={Flow Map Language Models: One-step Language Modeling via Continuous Denoising},
+      author={Chanhyuk Lee and Jaehoon Yoo and Manan Agarwal
+              and Sheel Shah and Jerry Huang
+              and Aditi Raghunathan and Seunghoon Hong
+              and Nicholas M. Boffi and Jinwoo Kim},
+      journal={arXiv preprint arXiv:2602.16813},
+      year={2026},
+  }
+  ```
+
+  ## Model Card Contact
+
+  Chanhyuk Lee (chanhyuk-lee@kaist.ac.kr)

checkpoints/quarantine/20260802T040621Z/official/duo/owt/README.md

@@ -0,0 +1,53 @@
+---
+datasets:
+- Skylion007/openwebtext
+language:
+- en
+library_name: transformers
+license: apache-2.0
+metrics:
+- perplexity
+pipeline_tag: text-generation
+---
+
+## Using DUO
+To use the pre-trained model for masked language modeling, use the following snippet:
+```python
+from transformers import AutoModelForMaskedLM, AutoTokenizer
+
+# See the `MDLM` collection page on the hub for list of available models.
+tokenizer = transformers.AutoTokenizer.from_pretrained('gpt2')
+model = AutoModelForMaskedLM.from_pretrained('s-sahoo/duo')
+```
+For a hands-on example, check out this [Colab notebook](https://colab.research.google.com/drive/1Sf7R-dqdR6gq-H8nyZ9E3ZkyvqMTqcwq?usp=sharing).
+For more information and implementation details, visit our github repository: [DUO](https://github.com/s-sahoo/duo) and project page: [Project Page](https://s-sahoo.com/duo)
+
+## Model Details
+The model, which has a context length of `1024` and is similar in size to GPT2-medium with approximately `130 million` non-embedding parameters,
+was trained for 1M steps on the OpenWebText corpus.
+
+For more details, please see our paper: [The Diffusion Duality](https://openreview.net/forum?id=CB0Ub2yXjC).
+
+
+
+## Citation
+
+<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
+Please cite our work using the bibtex below:
+
+**BibTeX:**
+
+```
+@inproceedings{
+sahoo2025the,
+title={The Diffusion Duality},
+author={Subham Sekhar Sahoo and Justin Deschenaux and Aaron Gokaslan and Guanghan Wang and Justin T Chiu and Volodymyr Kuleshov},
+booktitle={Forty-second International Conference on Machine Learning},
+year={2025},
+url={https://openreview.net/forum?id=9P9Y8FOSOk}
+}
+```
+
+
+## Model Card Contact
+Subham Sekhar Sahoo (ssahoo@cs.cornell.edu)

checkpoints/quarantine/20260802T040621Z/official/duo/owt/config.json

@@ -0,0 +1,23 @@
+{
+  "_name_or_path": "s-sahoo/duo",
+  "architectures": [
+    "DUO"
+  ],
+  "auto_map": {
+    "AutoConfig": "config.DUOConfig",
+    "AutoModelForMaskedLM": "model.DUO"
+  },
+  "causal": false,
+  "cond_dim": 128,
+  "dropout": 0.1,
+  "hidden_dim": 768,
+  "model_length": 1024,
+  "model_type": "DUO",
+  "n_blocks": 12,
+  "n_heads": 12,
+  "return_dict": false,
+  "torch_dtype": "float32",
+  "transformers_version": "4.38.2",
+  "var_min": true,
+  "vocab_size": 50258
+}

checkpoints/quarantine/20260802T040621Z/official/duo/owt/model.py

@@ -0,0 +1,601 @@
+import math
+import typing
+
+import einops
+import flash_attn
+import flash_attn.layers.rotary
+import huggingface_hub
+import omegaconf
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import transformers
+
+from .config import DUOConfig
+
+# Flags required to enable jit fusion kernels
+torch._C._jit_set_profiling_mode(False)
+torch._C._jit_set_profiling_executor(False)
+torch._C._jit_override_can_fuse_on_cpu(True)
+torch._C._jit_override_can_fuse_on_gpu(True)
+
+
+def bias_dropout_add_scale(
+    x: torch.Tensor,
+    bias: typing.Optional[torch.Tensor],
+    scale: torch.Tensor,
+    residual: typing.Optional[torch.Tensor],
+    prob: float,
+    training: bool) -> torch.Tensor:
+  if bias is not None:
+    out = scale * F.dropout(x + bias, p=prob, training=training)
+  else:
+    out = scale * F.dropout(x, p=prob, training=training)
+
+  if residual is not None:
+    out = residual + out
+  return out
+
+
+def get_bias_dropout_add_scale(training):
+  def _bias_dropout_add(x, bias, scale, residual, prob):
+    return bias_dropout_add_scale(
+      x, bias, scale, residual, prob, training)
+
+  return _bias_dropout_add
+
+
+# function overload
+def modulate(x: torch.Tensor,
+             shift: torch.Tensor,
+             scale: torch.Tensor) -> torch.Tensor:
+  return x * (1 + scale) + shift
+
+
+@torch.jit.script
+def bias_dropout_add_scale_fused_train(
+    x: torch.Tensor,
+    bias: typing.Optional[torch.Tensor],
+    scale: torch.Tensor,
+    residual: typing.Optional[torch.Tensor],
+    prob: float) -> torch.Tensor:
+  return bias_dropout_add_scale(
+    x, bias, scale, residual, prob, True)
+
+
+@torch.jit.script
+def bias_dropout_add_scale_fused_inference(
+    x: torch.Tensor,
+    bias: typing.Optional[torch.Tensor],
+    scale: torch.Tensor,
+    residual: typing.Optional[torch.Tensor],
+    prob: float) -> torch.Tensor:
+  return bias_dropout_add_scale(
+    x, bias, scale, residual, prob, False)
+
+
+@torch.jit.script
+def modulate_fused(x: torch.Tensor,
+                   shift: torch.Tensor,
+                   scale: torch.Tensor) -> torch.Tensor:
+  return modulate(x, shift, scale)
+
+
+class Rotary(torch.nn.Module):
+  def __init__(self, dim, base=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, seq_dim=1):
+    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.)
+      self.sin_cached[:,:,2,:,:].fill_(0.)
+
+    return self.cos_cached, self.sin_cached
+
+
+def rotate_half(x):
+  x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :]
+  return torch.cat((-x2, x1), dim=-1)
+
+
+def split_and_apply_rotary_pos_emb(qkv, rotary_cos_sin):
+  with torch.cuda.amp.autocast(enabled=False):
+    cos, sin = rotary_cos_sin
+    cos = cos.to(qkv.dtype)
+    sin = sin.to(qkv.dtype)
+    cos = cos[0,:,0,0,:cos.shape[-1]//2]
+    sin = sin[0,:,0,0,:sin.shape[-1]//2]
+    q, k, v = qkv.chunk(3, dim=2)
+    q = flash_attn.layers.rotary.apply_rotary_emb_torch(
+      q.squeeze(dim=2), cos, sin)
+    k = flash_attn.layers.rotary.apply_rotary_emb_torch(
+      k.squeeze(dim=2), cos, sin)
+    v = v.squeeze(dim=2)
+  return q, k, v
+
+
+def apply_rotary_pos_emb(qkv, cos, sin):
+  cos = cos[0,:,0,0,:cos.shape[-1]//2]
+  sin = sin[0,:,0,0,:sin.shape[-1]//2]
+  return flash_attn.layers.rotary.apply_rotary_emb_qkv_(qkv, cos, sin)
+
+
+def regular_attention_multi_headed(q, k, v):
+  # Assuming qkv is a tensor with shape [batch, seq_len, 3, num_heads, head_dim]
+  # where the 3 represents Q, K, V packed in that order
+  attention_output = F.scaled_dot_product_attention(
+    query=q.transpose(1, 2),
+    key=k.transpose(1, 2),
+    value=v.transpose(1, 2),
+    attn_mask=None,
+    dropout_p=0.0,
+    is_causal=False)
+  # [batch_size, seq_len, num_heads, head_dim]
+  attention_output = attention_output.transpose(1, 2)
+  return einops.rearrange(attention_output, 'b s h d -> b s (h d)')
+
+
+#################################################################################
+#                                  Layers                                       #
+#################################################################################
+class LayerNorm(nn.Module):
+  def __init__(self, dim):
+    super().__init__()
+    self.weight = nn.Parameter(torch.ones([dim]))
+    self.dim = dim
+  def forward(self, x):
+    with torch.cuda.amp.autocast(enabled=False):
+      x = F.layer_norm(x.float(), [self.dim])
+    return x * self.weight[None, None, :]
+
+
+def residual_linear(x, W, x_skip, residual_scale):
+  """x_skip + residual_scale * W @ x"""
+  dim_out, dim_in = W.shape[0], W.shape[1]
+  return torch.addmm(
+    x_skip.view(-1, dim_out),
+    x.view(-1, dim_in),
+    W.T,
+    alpha=residual_scale).view(*x.shape[:-1], dim_out)
+
+
+#################################################################################
+#               Embedding Layers for Timesteps and Class Labels                 #
+#################################################################################
+class TimestepEmbedder(nn.Module):
+  """
+  Embeds scalar timesteps into vector representations.
+  """
+  def __init__(self, hidden_size, frequency_embedding_size=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(t, dim, max_period=10000):
+    """
+    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.
+    """
+    # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
+    half = dim // 2
+    freqs = torch.exp(
+      - math.log(max_period)
+      * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device)
+      / half)
+    args = t[:, 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, t):
+    t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+    t_emb = self.mlp(t_freq)
+    return t_emb
+
+
+class LabelEmbedder(nn.Module):
+  """Embeds class labels into vector representations.
+  
+  Also handles label dropout for classifier-free guidance.
+  """
+  def __init__(self, num_classes, cond_size):
+    super().__init__()
+    self.embedding_table = nn.Embedding(num_classes + 1, cond_size)
+    self.num_classes = num_classes
+
+    # TODO think of initializing with 0.02 std deviation like in original DiT paper
+
+  def forward(self, labels):
+    embeddings = self.embedding_table(labels)
+    return embeddings
+    
+
+#################################################################################
+#                                 Core Model                                    #
+#################################################################################
+
+class DDiTBlockCausal(nn.Module):
+  def __init__(self, dim, n_heads, mlp_ratio=4, dropout=0.1):
+    super().__init__()
+    self.n_heads = n_heads
+
+    self.norm1 = LayerNorm(dim)
+    self.attn_qkv = nn.Linear(dim, 3 * dim, bias=False)
+    self.attn_out = nn.Linear(dim, dim, bias=False)
+    self.dropout1 = nn.Dropout(dropout)
+
+    self.norm2 = LayerNorm(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.dropout2 = nn.Dropout(dropout)
+    self.dropout = dropout
+
+  def _get_bias_dropout_scale(self):
+    if self.training:
+      return bias_dropout_add_scale_fused_train
+    else:
+      return bias_dropout_add_scale_fused_inference
+
+  def forward(self, x, rotary_cos_sin, **kwargs):
+    del kwargs
+    batch_size, seq_len = x.shape[0], x.shape[1]
+
+    bias_dropout_scale_fn = self._get_bias_dropout_scale()
+
+    # attention operation
+    x_skip = x
+    x = self.norm1(x)
+
+    qkv = self.attn_qkv(x)
+    qkv = einops.rearrange(
+      qkv,
+      'b s (three h d) -> b s three h d',
+      three=3,
+      h=self.n_heads)
+    with torch.cuda.amp.autocast(enabled=False):
+      cos, sin = rotary_cos_sin
+      qkv = apply_rotary_pos_emb(
+        qkv, cos.to(qkv.dtype), sin.to(qkv.dtype)
+      )
+    qkv = einops.rearrange(qkv, 'b s ... -> (b s) ...')
+    cu_seqlens = torch.arange(
+      0, (batch_size + 1) * seq_len,
+      step=seq_len, dtype=torch.int32, device=qkv.device)
+    x = flash_attn.flash_attn_interface.flash_attn_varlen_qkvpacked_func(
+      qkv, cu_seqlens, seq_len, 0.0, causal=True)
+
+    x = einops.rearrange(x, '(b s) h d -> b s (h d)',
+                         b=batch_size)
+
+    scale = torch.ones(1, device=x.device, dtype=x.dtype)
+    x = bias_dropout_scale_fn(
+      self.attn_out(x), None, scale, x_skip, self.dropout)
+
+    # mlp operation
+    x = bias_dropout_scale_fn(
+      self.mlp(self.norm2(x)), None, scale, x, self.dropout)
+    return x
+
+
+
+class DDiTBlock(nn.Module):
+  def __init__(self, dim, n_heads, adaLN,
+               cond_dim=None, mlp_ratio=4,
+               dropout=0.1):
+    super().__init__()
+    self.n_heads = n_heads
+    self.adaLN = adaLN
+
+    self.norm1 = LayerNorm(dim)
+    self.attn_qkv = nn.Linear(dim, 3 * dim, bias=False)
+    self.attn_out = nn.Linear(dim, dim, bias=False)
+    self.dropout1 = nn.Dropout(dropout)
+
+    self.norm2 = LayerNorm(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.dropout2 = nn.Dropout(dropout)
+    self.dropout = dropout
+
+    if self.adaLN:
+      self.adaLN_modulation = nn.Linear(cond_dim, 6 * dim)
+      self.adaLN_modulation.weight.data.zero_()
+      self.adaLN_modulation.bias.data.zero_()
+
+
+  def _get_bias_dropout_scale(self):
+    if self.training:
+      return bias_dropout_add_scale_fused_train
+    else:
+      return bias_dropout_add_scale_fused_inference
+
+
+  def forward(self, x, rotary_cos_sin, c=None):
+
+    bias_dropout_scale_fn = self._get_bias_dropout_scale()
+
+    x_skip = x
+    x = self.norm1(x)
+
+    if self.adaLN:
+      # self.adaLN_modulation(c): (128, 1536)
+      # self.adaLN_modulation(c)[:, None]: (128, 1, 1536)
+      # "" .chunk(6, dim=2) returns 6 tuples of shapes (128, 1, 256)
+      (shift_msa, scale_msa, gate_msa, shift_mlp,
+       scale_mlp, gate_mlp) = self.adaLN_modulation(c)[:, None].chunk(6, dim=2)
+      x = modulate_fused(x, shift_msa, scale_msa)
+
+    qkv = einops.rearrange(
+      self.attn_qkv(x),
+      'b s (three h d) -> b s three h d',
+      three=3,
+      h=self.n_heads)
+    q, k, v = split_and_apply_rotary_pos_emb(qkv, rotary_cos_sin)
+    
+    x = regular_attention_multi_headed(q, k, v)
+
+    if self.adaLN:
+      x = bias_dropout_scale_fn(self.attn_out(x),
+                                None,
+                                gate_msa,
+                                x_skip,
+                                self.dropout)
+      x = bias_dropout_scale_fn(
+        self.mlp(modulate_fused(
+          self.norm2(x), shift_mlp, scale_mlp)),
+        None, gate_mlp, x, self.dropout)
+    else:
+      scale = torch.ones(1, device=x.device, dtype=x.dtype)
+      x = bias_dropout_scale_fn(
+        self.attn_out(x), None, scale, x_skip, self.dropout)
+      x = bias_dropout_scale_fn(
+        self.mlp(self.norm2(x)), None, scale, x, self.dropout)
+    return x
+
+
+class EmbeddingLayer(nn.Module):
+  def __init__(self, dim, vocab_dim):
+    super().__init__()
+    self.embedding = nn.Parameter(torch.empty((vocab_dim, dim)))
+    torch.nn.init.kaiming_uniform_(self.embedding, a=math.sqrt(5))
+
+  def forward(self, x):
+    if x.ndim == 2:
+      return self.embedding[x]
+    assert x.ndim == 3
+    return torch.einsum(
+      "blv,ve->ble",
+      torch.nn.functional.softmax(x, dim=-1).float(),
+      self.embedding.float()).to(x.dtype)
+
+
+class DDiTFinalLayer(nn.Module):
+  def __init__(self, hidden_size, out_channels, cond_dim,
+               adaLN):
+    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 = adaLN
+    if self.adaLN:
+      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, c):
+    x = self.norm_final(x)
+    if self.adaLN:
+      shift, scale = self.adaLN_modulation(c)[:, None].chunk(2, dim=2)
+      x = modulate_fused(x, shift, scale)
+    x = self.linear(x)
+    return x
+
+
+class DIT(nn.Module, huggingface_hub.PyTorchModelHubMixin):
+  def __init__(self, config, vocab_size: int):
+    super().__init__()
+    if type(config) == dict:
+      config = omegaconf.OmegaConf.create(config)
+    self.causal = config.algo.causal_attention
+    self.adaLN = not self.causal
+    self.config = config
+    self.vocab_size = vocab_size
+    dim = config.model.hidden_size
+    cond_dim = config.model.cond_dim
+    self.vocab_embed = EmbeddingLayer(dim, vocab_size)
+    if not self.causal:
+      self.sigma_map = TimestepEmbedder(cond_dim)
+    self.rotary_emb = Rotary(dim // config.model.n_heads)
+
+    blocks = []
+    for _ in range(config.model.n_blocks):
+      if self.causal:
+        block = DDiTBlockCausal(
+          dim=dim,
+          n_heads=config.model.n_heads,
+          dropout=config.model.dropout)
+      else:
+        block = DDiTBlock(
+          dim=dim,
+          n_heads=config.model.n_heads,
+          cond_dim=cond_dim,
+          adaLN=self.adaLN,
+          dropout=config.model.dropout)
+      blocks.append(block)
+    self.blocks = nn.ModuleList(blocks)
+
+    self.output_layer = DDiTFinalLayer(
+      hidden_size=dim,
+      out_channels=vocab_size,
+      cond_dim=cond_dim,
+      adaLN=self.adaLN)
+    self.scale_by_sigma = config.model.scale_by_sigma
+
+  def _get_bias_dropout_scale(self):
+    if self.training:
+      return bias_dropout_add_scale_fused_train
+    else:
+      return  bias_dropout_add_scale_fused_inference
+
+  def forward(self, x, sigma):
+    x = self.vocab_embed(x)
+    if self.causal:
+      t_cond = None
+    else:
+      t_cond = F.silu(self.sigma_map(sigma))
+
+    rotary_cos_sin = self.rotary_emb(x).to(x.device)
+
+    with torch.cuda.amp.autocast(dtype=torch.bfloat16):
+      for i in range(len(self.blocks)):
+        x = self.blocks[i](x, rotary_cos_sin, c=t_cond)
+      x = self.output_layer(x, c=t_cond)
+
+    return x
+
+
+
+class HFDIT(torch.nn.Module):
+  def __init__(self, config):
+    super().__init__()
+    self.causal = config.causal
+    self.adaLN = not self.causal
+    self.vocab_size = config.vocab_size
+    dim = config.hidden_dim
+    cond_dim = config.cond_dim
+    self.vocab_embed = EmbeddingLayer(dim, self.vocab_size)
+    if not self.causal:
+      self.sigma_map = TimestepEmbedder(cond_dim)
+    self.rotary_emb = Rotary(dim // config.n_heads)
+
+    blocks = []
+    for _ in range(config.n_blocks):
+      if self.causal:
+        block = DDiTBlockCausal(
+          dim=dim,
+          n_heads=config.n_heads,
+          dropout=config.dropout)
+      else:
+        block = DDiTBlock(
+          dim=dim,
+          n_heads=config.n_heads,
+          cond_dim=cond_dim,
+          adaLN=self.adaLN,
+          dropout=config.dropout)
+      blocks.append(block)
+    self.blocks = torch.nn.ModuleList(blocks)
+
+    self.output_layer = DDiTFinalLayer(
+      hidden_size=dim,
+      out_channels=self.vocab_size,
+      cond_dim=cond_dim,
+      adaLN=self.adaLN)
+
+  def _get_bias_dropout_scale(self):
+    if self.training:
+      return bias_dropout_add_scale_fused_train
+    else:
+      return  bias_dropout_add_scale_fused_inference
+
+  def forward(self, x, sigma, output_hidden_states=False):
+    all_hidden_states = []
+    x = self.vocab_embed(x)
+    if output_hidden_states:
+      all_hidden_states.append(x)
+    if self.causal:
+      t_cond = None
+    else:
+      t_cond = F.silu(self.sigma_map(sigma))
+
+    rotary_cos_sin = self.rotary_emb(x)
+    with torch.cuda.amp.autocast(
+      dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float32):
+      for i in range(len(self.blocks)):
+        x = self.blocks[i](x, rotary_cos_sin, c=t_cond)
+        if output_hidden_states:
+          all_hidden_states.append(x)
+      x = self.output_layer(x, c=t_cond)
+    return x, all_hidden_states
+
+
+class DUO(transformers.PreTrainedModel):
+  """HF-compatible model."""
+  config_class = DUOConfig
+  base_model_prefix = 'duo'
+
+  def __init__(self, config: DUOConfig):
+    super().__init__(config)
+    self.config = config
+    self.backbone = HFDIT(config)
+
+  def reset_kv_cache(self):
+    for block in self.backbone.blocks:
+      block.kv_cache = None
+
+  def forward(
+      self,
+      input_ids: torch.LongTensor = None,
+      timesteps: torch.FloatTensor = None,
+      output_hidden_states: typing.Optional[bool] = None,
+      return_dict: typing.Optional[bool] = None,
+  ) -> typing.Union[
+    torch.Tensor, typing.Tuple,
+    transformers.modeling_outputs.MaskedLMOutput]:
+    """HF-compatible forward method."""
+    output_hidden_states = (
+      output_hidden_states
+      if output_hidden_states is not None
+      else self.config.output_hidden_states
+    )
+    return_dict = return_dict \
+      if return_dict is not None \
+      else self.config.use_return_dict
+    
+    logits, all_hidden_states = self.backbone(
+      x=input_ids,
+      sigma=timesteps,
+      output_hidden_states=output_hidden_states,
+    )
+    if return_dict:
+      return transformers.modeling_outputs.MaskedLMOutput(
+        logits=logits,
+        hidden_states=all_hidden_states if output_hidden_states else None,
+        loss=None
+      )
+    elif output_hidden_states:
+      return logits, all_hidden_states
+    else:
+      return logits

checkpoints/quarantine/20260802T040629Z/official/duo_dcd/owt/README.md

@@ -0,0 +1,52 @@
+---
+datasets:
+- Skylion007/openwebtext
+language:
+- en
+library_name: transformers
+license: apache-2.0
+metrics:
+- perplexity
+pipeline_tag: text-generation
+---
+
+## Using DUO
+To use the pre-trained model for masked language modeling, use the following snippet:
+```python
+from transformers import AutoModelForMaskedLM, AutoTokenizer
+
+# See the `DUO` collection page on the hub for list of available models.
+tokenizer = transformers.AutoTokenizer.from_pretrained('gpt2')
+model = AutoModelForMaskedLM.from_pretrained('s-sahoo/duo-distilled')
+```
+For a hands-on example, check out this [Colab notebook](https://colab.research.google.com/drive/1Sf7R-dqdR6gq-H8nyZ9E3ZkyvqMTqcwq?usp=sharing).
+For more information and implementation details, visit our github repository: [DUO](https://github.com/s-sahoo/duo)
+
+## Model Details
+The model, which has a context length of `1024` and is similar in size to GPT2-medium with approximately `130 million` non-embedding parameters,
+was trained for 1M steps on the OpenWebText corpus.
+
+For more details, please see our paper: [The Diffusion Duality](https://huggingface.co/papers/2506.10892).
+
+Project page: https://s-sahoo.com/duo
+
+## Citation
+
+<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
+Please cite our work using the bibtex below:
+
+**BibTeX:**
+
+```
+@inproceedings{
+sahoo2025the,
+title={The Diffusion Duality},
+author={Subham Sekhar Sahoo and Justin Deschenaux and Aaron Gokaslan and Guanghan Wang and Justin T Chiu and Volodymyr Kuleshov},
+booktitle={Forty-second International Conference on Machine Learning},
+year={2025},
+url={https://openreview.net/forum?id=9P9Y8FOSOk}
+}
+```
+
+## Model Card Contact
+Subham Sekhar Sahoo (ssahoo@cs.cornell.edu)

checkpoints/quarantine/20260802T040629Z/official/duo_dcd/owt/config.json

@@ -0,0 +1,23 @@
+{
+  "_name_or_path": "s-sahoo/duo-distilled",
+  "architectures": [
+    "DUO"
+  ],
+  "auto_map": {
+    "AutoConfig": "config.DUOConfig",
+    "AutoModelForMaskedLM": "model.DUO"
+  },
+  "causal": false,
+  "cond_dim": 128,
+  "dropout": 0.1,
+  "hidden_dim": 768,
+  "model_length": 1024,
+  "model_type": "DUO",
+  "n_blocks": 12,
+  "n_heads": 12,
+  "return_dict": false,
+  "torch_dtype": "float32",
+  "transformers_version": "4.38.2",
+  "var_min": true,
+  "vocab_size": 50258
+}

checkpoints/quarantine/20260802T040629Z/official/duo_dcd/owt/model.py

@@ -0,0 +1,601 @@
+import math
+import typing
+
+import einops
+import flash_attn
+import flash_attn.layers.rotary
+import huggingface_hub
+import omegaconf
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import transformers
+
+from .config import DUOConfig
+
+# Flags required to enable jit fusion kernels
+torch._C._jit_set_profiling_mode(False)
+torch._C._jit_set_profiling_executor(False)
+torch._C._jit_override_can_fuse_on_cpu(True)
+torch._C._jit_override_can_fuse_on_gpu(True)
+
+
+def bias_dropout_add_scale(
+    x: torch.Tensor,
+    bias: typing.Optional[torch.Tensor],
+    scale: torch.Tensor,
+    residual: typing.Optional[torch.Tensor],
+    prob: float,
+    training: bool) -> torch.Tensor:
+  if bias is not None:
+    out = scale * F.dropout(x + bias, p=prob, training=training)
+  else:
+    out = scale * F.dropout(x, p=prob, training=training)
+
+  if residual is not None:
+    out = residual + out
+  return out
+
+
+def get_bias_dropout_add_scale(training):
+  def _bias_dropout_add(x, bias, scale, residual, prob):
+    return bias_dropout_add_scale(
+      x, bias, scale, residual, prob, training)
+
+  return _bias_dropout_add
+
+
+# function overload
+def modulate(x: torch.Tensor,
+             shift: torch.Tensor,
+             scale: torch.Tensor) -> torch.Tensor:
+  return x * (1 + scale) + shift
+
+
+@torch.jit.script
+def bias_dropout_add_scale_fused_train(
+    x: torch.Tensor,
+    bias: typing.Optional[torch.Tensor],
+    scale: torch.Tensor,
+    residual: typing.Optional[torch.Tensor],
+    prob: float) -> torch.Tensor:
+  return bias_dropout_add_scale(
+    x, bias, scale, residual, prob, True)
+
+
+@torch.jit.script
+def bias_dropout_add_scale_fused_inference(
+    x: torch.Tensor,
+    bias: typing.Optional[torch.Tensor],
+    scale: torch.Tensor,
+    residual: typing.Optional[torch.Tensor],
+    prob: float) -> torch.Tensor:
+  return bias_dropout_add_scale(
+    x, bias, scale, residual, prob, False)
+
+
+@torch.jit.script
+def modulate_fused(x: torch.Tensor,
+                   shift: torch.Tensor,
+                   scale: torch.Tensor) -> torch.Tensor:
+  return modulate(x, shift, scale)
+
+
+class Rotary(torch.nn.Module):
+  def __init__(self, dim, base=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, seq_dim=1):
+    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.)
+      self.sin_cached[:,:,2,:,:].fill_(0.)
+
+    return self.cos_cached, self.sin_cached
+
+
+def rotate_half(x):
+  x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :]
+  return torch.cat((-x2, x1), dim=-1)
+
+
+def split_and_apply_rotary_pos_emb(qkv, rotary_cos_sin):
+  with torch.cuda.amp.autocast(enabled=False):
+    cos, sin = rotary_cos_sin
+    cos = cos.to(qkv.dtype)
+    sin = sin.to(qkv.dtype)
+    cos = cos[0,:,0,0,:cos.shape[-1]//2]
+    sin = sin[0,:,0,0,:sin.shape[-1]//2]
+    q, k, v = qkv.chunk(3, dim=2)
+    q = flash_attn.layers.rotary.apply_rotary_emb_torch(
+      q.squeeze(dim=2), cos, sin)
+    k = flash_attn.layers.rotary.apply_rotary_emb_torch(
+      k.squeeze(dim=2), cos, sin)
+    v = v.squeeze(dim=2)
+  return q, k, v
+
+
+def apply_rotary_pos_emb(qkv, cos, sin):
+  cos = cos[0,:,0,0,:cos.shape[-1]//2]
+  sin = sin[0,:,0,0,:sin.shape[-1]//2]
+  return flash_attn.layers.rotary.apply_rotary_emb_qkv_(qkv, cos, sin)
+
+
+def regular_attention_multi_headed(q, k, v):
+  # Assuming qkv is a tensor with shape [batch, seq_len, 3, num_heads, head_dim]
+  # where the 3 represents Q, K, V packed in that order
+  attention_output = F.scaled_dot_product_attention(
+    query=q.transpose(1, 2),
+    key=k.transpose(1, 2),
+    value=v.transpose(1, 2),
+    attn_mask=None,
+    dropout_p=0.0,
+    is_causal=False)
+  # [batch_size, seq_len, num_heads, head_dim]
+  attention_output = attention_output.transpose(1, 2)
+  return einops.rearrange(attention_output, 'b s h d -> b s (h d)')
+
+
+#################################################################################
+#                                  Layers                                       #
+#################################################################################
+class LayerNorm(nn.Module):
+  def __init__(self, dim):
+    super().__init__()
+    self.weight = nn.Parameter(torch.ones([dim]))
+    self.dim = dim
+  def forward(self, x):
+    with torch.cuda.amp.autocast(enabled=False):
+      x = F.layer_norm(x.float(), [self.dim])
+    return x * self.weight[None, None, :]
+
+
+def residual_linear(x, W, x_skip, residual_scale):
+  """x_skip + residual_scale * W @ x"""
+  dim_out, dim_in = W.shape[0], W.shape[1]
+  return torch.addmm(
+    x_skip.view(-1, dim_out),
+    x.view(-1, dim_in),
+    W.T,
+    alpha=residual_scale).view(*x.shape[:-1], dim_out)
+
+
+#################################################################################
+#               Embedding Layers for Timesteps and Class Labels                 #
+#################################################################################
+class TimestepEmbedder(nn.Module):
+  """
+  Embeds scalar timesteps into vector representations.
+  """
+  def __init__(self, hidden_size, frequency_embedding_size=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(t, dim, max_period=10000):
+    """
+    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.
+    """
+    # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
+    half = dim // 2
+    freqs = torch.exp(
+      - math.log(max_period)
+      * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device)
+      / half)
+    args = t[:, 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, t):
+    t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+    t_emb = self.mlp(t_freq)
+    return t_emb
+
+
+class LabelEmbedder(nn.Module):
+  """Embeds class labels into vector representations.
+  
+  Also handles label dropout for classifier-free guidance.
+  """
+  def __init__(self, num_classes, cond_size):
+    super().__init__()
+    self.embedding_table = nn.Embedding(num_classes + 1, cond_size)
+    self.num_classes = num_classes
+
+    # TODO think of initializing with 0.02 std deviation like in original DiT paper
+
+  def forward(self, labels):
+    embeddings = self.embedding_table(labels)
+    return embeddings
+    
+
+#################################################################################
+#                                 Core Model                                    #
+#################################################################################
+
+class DDiTBlockCausal(nn.Module):
+  def __init__(self, dim, n_heads, mlp_ratio=4, dropout=0.1):
+    super().__init__()
+    self.n_heads = n_heads
+
+    self.norm1 = LayerNorm(dim)
+    self.attn_qkv = nn.Linear(dim, 3 * dim, bias=False)
+    self.attn_out = nn.Linear(dim, dim, bias=False)
+    self.dropout1 = nn.Dropout(dropout)
+
+    self.norm2 = LayerNorm(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.dropout2 = nn.Dropout(dropout)
+    self.dropout = dropout
+
+  def _get_bias_dropout_scale(self):
+    if self.training:
+      return bias_dropout_add_scale_fused_train
+    else:
+      return bias_dropout_add_scale_fused_inference
+
+  def forward(self, x, rotary_cos_sin, **kwargs):
+    del kwargs
+    batch_size, seq_len = x.shape[0], x.shape[1]
+
+    bias_dropout_scale_fn = self._get_bias_dropout_scale()
+
+    # attention operation
+    x_skip = x
+    x = self.norm1(x)
+
+    qkv = self.attn_qkv(x)
+    qkv = einops.rearrange(
+      qkv,
+      'b s (three h d) -> b s three h d',
+      three=3,
+      h=self.n_heads)
+    with torch.cuda.amp.autocast(enabled=False):
+      cos, sin = rotary_cos_sin
+      qkv = apply_rotary_pos_emb(
+        qkv, cos.to(qkv.dtype), sin.to(qkv.dtype)
+      )
+    qkv = einops.rearrange(qkv, 'b s ... -> (b s) ...')
+    cu_seqlens = torch.arange(
+      0, (batch_size + 1) * seq_len,
+      step=seq_len, dtype=torch.int32, device=qkv.device)
+    x = flash_attn.flash_attn_interface.flash_attn_varlen_qkvpacked_func(
+      qkv, cu_seqlens, seq_len, 0.0, causal=True)
+
+    x = einops.rearrange(x, '(b s) h d -> b s (h d)',
+                         b=batch_size)
+
+    scale = torch.ones(1, device=x.device, dtype=x.dtype)
+    x = bias_dropout_scale_fn(
+      self.attn_out(x), None, scale, x_skip, self.dropout)
+
+    # mlp operation
+    x = bias_dropout_scale_fn(
+      self.mlp(self.norm2(x)), None, scale, x, self.dropout)
+    return x
+
+
+
+class DDiTBlock(nn.Module):
+  def __init__(self, dim, n_heads, adaLN,
+               cond_dim=None, mlp_ratio=4,
+               dropout=0.1):
+    super().__init__()
+    self.n_heads = n_heads
+    self.adaLN = adaLN
+
+    self.norm1 = LayerNorm(dim)
+    self.attn_qkv = nn.Linear(dim, 3 * dim, bias=False)
+    self.attn_out = nn.Linear(dim, dim, bias=False)
+    self.dropout1 = nn.Dropout(dropout)
+
+    self.norm2 = LayerNorm(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.dropout2 = nn.Dropout(dropout)
+    self.dropout = dropout
+
+    if self.adaLN:
+      self.adaLN_modulation = nn.Linear(cond_dim, 6 * dim)
+      self.adaLN_modulation.weight.data.zero_()
+      self.adaLN_modulation.bias.data.zero_()
+
+
+  def _get_bias_dropout_scale(self):
+    if self.training:
+      return bias_dropout_add_scale_fused_train
+    else:
+      return bias_dropout_add_scale_fused_inference
+
+
+  def forward(self, x, rotary_cos_sin, c=None):
+
+    bias_dropout_scale_fn = self._get_bias_dropout_scale()
+
+    x_skip = x
+    x = self.norm1(x)
+
+    if self.adaLN:
+      # self.adaLN_modulation(c): (128, 1536)
+      # self.adaLN_modulation(c)[:, None]: (128, 1, 1536)
+      # "" .chunk(6, dim=2) returns 6 tuples of shapes (128, 1, 256)
+      (shift_msa, scale_msa, gate_msa, shift_mlp,
+       scale_mlp, gate_mlp) = self.adaLN_modulation(c)[:, None].chunk(6, dim=2)
+      x = modulate_fused(x, shift_msa, scale_msa)
+
+    qkv = einops.rearrange(
+      self.attn_qkv(x),
+      'b s (three h d) -> b s three h d',
+      three=3,
+      h=self.n_heads)
+    q, k, v = split_and_apply_rotary_pos_emb(qkv, rotary_cos_sin)
+    
+    x = regular_attention_multi_headed(q, k, v)
+
+    if self.adaLN:
+      x = bias_dropout_scale_fn(self.attn_out(x),
+                                None,
+                                gate_msa,
+                                x_skip,
+                                self.dropout)
+      x = bias_dropout_scale_fn(
+        self.mlp(modulate_fused(
+          self.norm2(x), shift_mlp, scale_mlp)),
+        None, gate_mlp, x, self.dropout)
+    else:
+      scale = torch.ones(1, device=x.device, dtype=x.dtype)
+      x = bias_dropout_scale_fn(
+        self.attn_out(x), None, scale, x_skip, self.dropout)
+      x = bias_dropout_scale_fn(
+        self.mlp(self.norm2(x)), None, scale, x, self.dropout)
+    return x
+
+
+class EmbeddingLayer(nn.Module):
+  def __init__(self, dim, vocab_dim):
+    super().__init__()
+    self.embedding = nn.Parameter(torch.empty((vocab_dim, dim)))
+    torch.nn.init.kaiming_uniform_(self.embedding, a=math.sqrt(5))
+
+  def forward(self, x):
+    if x.ndim == 2:
+      return self.embedding[x]
+    assert x.ndim == 3
+    return torch.einsum(
+      "blv,ve->ble",
+      torch.nn.functional.softmax(x, dim=-1).float(),
+      self.embedding.float()).to(x.dtype)
+
+
+class DDiTFinalLayer(nn.Module):
+  def __init__(self, hidden_size, out_channels, cond_dim,
+               adaLN):
+    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 = adaLN
+    if self.adaLN:
+      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, c):
+    x = self.norm_final(x)
+    if self.adaLN:
+      shift, scale = self.adaLN_modulation(c)[:, None].chunk(2, dim=2)
+      x = modulate_fused(x, shift, scale)
+    x = self.linear(x)
+    return x
+
+
+class DIT(nn.Module, huggingface_hub.PyTorchModelHubMixin):
+  def __init__(self, config, vocab_size: int):
+    super().__init__()
+    if type(config) == dict:
+      config = omegaconf.OmegaConf.create(config)
+    self.causal = config.algo.causal_attention
+    self.adaLN = not self.causal
+    self.config = config
+    self.vocab_size = vocab_size
+    dim = config.model.hidden_size
+    cond_dim = config.model.cond_dim
+    self.vocab_embed = EmbeddingLayer(dim, vocab_size)
+    if not self.causal:
+      self.sigma_map = TimestepEmbedder(cond_dim)
+    self.rotary_emb = Rotary(dim // config.model.n_heads)
+
+    blocks = []
+    for _ in range(config.model.n_blocks):
+      if self.causal:
+        block = DDiTBlockCausal(
+          dim=dim,
+          n_heads=config.model.n_heads,
+          dropout=config.model.dropout)
+      else:
+        block = DDiTBlock(
+          dim=dim,
+          n_heads=config.model.n_heads,
+          cond_dim=cond_dim,
+          adaLN=self.adaLN,
+          dropout=config.model.dropout)
+      blocks.append(block)
+    self.blocks = nn.ModuleList(blocks)
+
+    self.output_layer = DDiTFinalLayer(
+      hidden_size=dim,
+      out_channels=vocab_size,
+      cond_dim=cond_dim,
+      adaLN=self.adaLN)
+    self.scale_by_sigma = config.model.scale_by_sigma
+
+  def _get_bias_dropout_scale(self):
+    if self.training:
+      return bias_dropout_add_scale_fused_train
+    else:
+      return  bias_dropout_add_scale_fused_inference
+
+  def forward(self, x, sigma):
+    x = self.vocab_embed(x)
+    if self.causal:
+      t_cond = None
+    else:
+      t_cond = F.silu(self.sigma_map(sigma))
+
+    rotary_cos_sin = self.rotary_emb(x).to(x.device)
+
+    with torch.cuda.amp.autocast(dtype=torch.bfloat16):
+      for i in range(len(self.blocks)):
+        x = self.blocks[i](x, rotary_cos_sin, c=t_cond)
+      x = self.output_layer(x, c=t_cond)
+
+    return x
+
+
+
+class HFDIT(torch.nn.Module):
+  def __init__(self, config):
+    super().__init__()
+    self.causal = config.causal
+    self.adaLN = not self.causal
+    self.vocab_size = config.vocab_size
+    dim = config.hidden_dim
+    cond_dim = config.cond_dim
+    self.vocab_embed = EmbeddingLayer(dim, self.vocab_size)
+    if not self.causal:
+      self.sigma_map = TimestepEmbedder(cond_dim)
+    self.rotary_emb = Rotary(dim // config.n_heads)
+
+    blocks = []
+    for _ in range(config.n_blocks):
+      if self.causal:
+        block = DDiTBlockCausal(
+          dim=dim,
+          n_heads=config.n_heads,
+          dropout=config.dropout)
+      else:
+        block = DDiTBlock(
+          dim=dim,
+          n_heads=config.n_heads,
+          cond_dim=cond_dim,
+          adaLN=self.adaLN,
+          dropout=config.dropout)
+      blocks.append(block)
+    self.blocks = torch.nn.ModuleList(blocks)
+
+    self.output_layer = DDiTFinalLayer(
+      hidden_size=dim,
+      out_channels=self.vocab_size,
+      cond_dim=cond_dim,
+      adaLN=self.adaLN)
+
+  def _get_bias_dropout_scale(self):
+    if self.training:
+      return bias_dropout_add_scale_fused_train
+    else:
+      return  bias_dropout_add_scale_fused_inference
+
+  def forward(self, x, sigma, output_hidden_states=False):
+    all_hidden_states = []
+    x = self.vocab_embed(x)
+    if output_hidden_states:
+      all_hidden_states.append(x)
+    if self.causal:
+      t_cond = None
+    else:
+      t_cond = F.silu(self.sigma_map(sigma))
+
+    rotary_cos_sin = self.rotary_emb(x)
+    with torch.cuda.amp.autocast(
+      dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float32):
+      for i in range(len(self.blocks)):
+        x = self.blocks[i](x, rotary_cos_sin, c=t_cond)
+        if output_hidden_states:
+          all_hidden_states.append(x)
+      x = self.output_layer(x, c=t_cond)
+    return x, all_hidden_states
+
+
+class DUO(transformers.PreTrainedModel):
+  """HF-compatible model."""
+  config_class = DUOConfig
+  base_model_prefix = 'duo'
+
+  def __init__(self, config: DUOConfig):
+    super().__init__(config)
+    self.config = config
+    self.backbone = HFDIT(config)
+
+  def reset_kv_cache(self):
+    for block in self.backbone.blocks:
+      block.kv_cache = None
+
+  def forward(
+      self,
+      input_ids: torch.LongTensor = None,
+      timesteps: torch.FloatTensor = None,
+      output_hidden_states: typing.Optional[bool] = None,
+      return_dict: typing.Optional[bool] = None,
+  ) -> typing.Union[
+    torch.Tensor, typing.Tuple,
+    transformers.modeling_outputs.MaskedLMOutput]:
+    """HF-compatible forward method."""
+    output_hidden_states = (
+      output_hidden_states
+      if output_hidden_states is not None
+      else self.config.output_hidden_states
+    )
+    return_dict = return_dict \
+      if return_dict is not None \
+      else self.config.use_return_dict
+    
+    logits, all_hidden_states = self.backbone(
+      x=input_ids,
+      sigma=timesteps,
+      output_hidden_states=output_hidden_states,
+    )
+    if return_dict:
+      return transformers.modeling_outputs.MaskedLMOutput(
+        logits=logits,
+        hidden_states=all_hidden_states if output_hidden_states else None,
+        loss=None
+      )
+    elif output_hidden_states:
+      return logits, all_hidden_states
+    else:
+      return logits

checkpoints/quarantine/20260802T040645Z/official/mdlm_sdtt/owt/config.json

@@ -0,0 +1,207 @@
+{
+  "mode": "train",
+  "time_conditioning": false,
+  "T": 0,
+  "compile": false,
+  "eval_every": 5000,
+  "data_preprocess": {
+    "data_cache": "REDACTED",
+    "min_seq_len": -1,
+    "seq_len": "${model.length}",
+    "group_text": true,
+    "remove_text": true,
+    "num_seqs": -1,
+    "add_bos": false,
+    "add_eos": true,
+    "legacy_start_end_bos": true
+  },
+  "loader": {
+    "global_batch_size": 32,
+    "eval_global_batch_size": "${.global_batch_size}",
+    "batch_size": "${div_up:${.global_batch_size}, ${eval:${trainer.devices} * ${trainer.num_nodes}}}",
+    "eval_batch_size": "${div_up:${.eval_global_batch_size}, ${eval:${trainer.devices} * ${trainer.num_nodes}}}",
+    "num_workers": 32,
+    "pin_memory": true,
+    "persistent_workers": true
+  },
+  "training": {
+    "ema": 0.9999,
+    "antithetic_sampling": true,
+    "importance_sampling": false,
+    "sampling_eps": 1e-05,
+    "change_of_variables": false
+  },
+  "tokenizer": {
+    "name": "gpt2"
+  },
+  "optim": {
+    "name": "adamw",
+    "weight_decay": 0.0,
+    "lr": 0.0003,
+    "beta1": 0.9,
+    "beta2": 0.999,
+    "eps": 1e-08
+  },
+  "trainer": {
+    "_target_": "lightning.Trainer",
+    "accelerator": "cuda",
+    "num_nodes": 1,
+    "devices": 1,
+    "accumulate_grad_batches": "${div_up:${loader.global_batch_size}, ${eval:${trainer.devices} * ${loader.batch_size} * ${trainer.num_nodes}}}",
+    "gradient_clip_val": 1.0,
+    "precision": "bf16-mixed",
+    "num_sanity_val_steps": 2,
+    "max_steps": 1000000,
+    "log_every_n_steps": 10,
+    "limit_train_batches": 1.0,
+    "limit_val_batches": 1.0,
+    "val_check_interval": "${eval:${trainer.accumulate_grad_batches} * ${eval_every}}",
+    "check_val_every_n_epoch": null,
+    "benchmark": true
+  },
+  "wandb": {
+    "project": "text-diffusion",
+    "notes": "",
+    "group": null,
+    "job_type": null,
+    "name": null,
+    "tags": [
+      "${noise.type}",
+      "${data.train}",
+      "${data.valid}"
+    ]
+  },
+  "checkpointing": {
+    "save_dir": "${cwd:}",
+    "resume_from_ckpt": true,
+    "resume_ckpt_path": "/claire-rcp-scratch/home/deschena/fast-discrete-diff/dev/hf_downloads/baselines/kld/baselines/bwd_kl_double_dt_10k_lr_2x_sm/70000.ckpt"
+  },
+  "eval": {
+    "valid": {
+      "n_samples": 16,
+      "num_steps": 128
+    },
+    "ppl_with_ar": {
+      "run": true,
+      "model": "gpt2-large",
+      "batch_size": 8
+    },
+    "mauve": {
+      "run": false,
+      "num_rounds": 5,
+      "max_num_tokens": 100,
+      "batch_size": 8,
+      "scaling_factor": 5
+    },
+    "lambada_openai": {
+      "run": false,
+      "batch_size": 32,
+      "num_samples": 32,
+      "from_ema": true,
+      "add_bos": true,
+      "add_eos": true
+    }
+  },
+  "callbacks": {
+    "checkpoint_every_n_steps": {
+      "_target_": "lightning.pytorch.callbacks.ModelCheckpoint",
+      "save_top_k": -1,
+      "save_last": true,
+      "dirpath": "${checkpointing.save_dir}/checkpoints",
+      "verbose": true,
+      "auto_insert_metric_name": false,
+      "every_n_train_steps": 10000
+    },
+    "learning_rate_monitor": {
+      "_target_": "lightning.pytorch.callbacks.LearningRateMonitor",
+      "logging_interval": "step"
+    },
+    "text_throughput": {
+      "_target_": "callbacks.text_throughput.TextThroughputMonitor"
+    },
+    "grad_norm": {
+      "_target_": "callbacks.gradient_norm.GradNormCallback"
+    }
+  },
+  "data": {
+    "train": "openwebtext-train",
+    "valid": "openwebtext-valid",
+    "test": "openwebtext-test",
+    "tokenizer_name_or_path": "gpt2",
+    "wrap": true,
+    "streaming": false
+  },
+  "model": {
+    "name": "small",
+    "type": "ddit-orig",
+    "hidden_size": 768,
+    "cond_dim": 128,
+    "length": 1024,
+    "n_blocks": 12,
+    "n_heads": 12,
+    "scale_by_sigma": true,
+    "dropout": 0.1,
+    "tie_word_embeddings": false
+  },
+  "strategy": {
+    "_target_": "lightning.pytorch.strategies.DDPStrategy",
+    "find_unused_parameters": false,
+    "process_group_backend": "nccl"
+  },
+  "noise": {
+    "type": "loglinear",
+    "sigma_min": 0.0001,
+    "sigma_max": 20
+  },
+  "lr_scheduler": {
+    "_target_": "transformers.get_constant_schedule_with_warmup",
+    "num_warmup_steps": 2500
+  },
+  "parameterization": {
+    "name": "multi-round-sdtt",
+    "log_loss_buckets": -1,
+    "checkpoint_path": "kuleshov-group/mdlm-owt",
+    "start_from_hf": false,
+    "distill_mode": "kl-bwd",
+    "num_distill_steps": 2,
+    "min_num_sampling_steps": 8,
+    "grow_dt_every": 10000,
+    "orig_num_sampling_steps": 1024,
+    "sampling_mode": "ancestral",
+    "loss_precision": null,
+    "reset_optimizer_on_growth": false,
+    "use_ema_on_growth": false,
+    "sampling": {
+      "uncond": {
+        "run": false,
+        "num_samples": 1024,
+        "batch_size": 32,
+        "from_ema": true,
+        "num_steps": 1024,
+        "seq_len": 1024,
+        "sampler": "ancestral",
+        "cache_preds": false,
+        "add_bos": false,
+        "add_eos": false
+      },
+      "cond_prefix": {
+        "run": false,
+        "num_samples": 1024,
+        "batch_size": 32,
+        "from_ema": true,
+        "dataset": "webtext",
+        "seq_len": 100,
+        "prefix_len": 50,
+        "num_cont_per_prefix": 5,
+        "min_seq_len": 1024,
+        "num_steps": 1024,
+        "sampler": "ancestral",
+        "cache_preds": false,
+        "add_bos": false,
+        "add_eos": false
+      }
+    }
+  },
+  "repo": "jdeschena/sdtt",
+  "revision": "baselines_kld"
+}

checkpoints/quarantine/20260818T074854Z/official/duo_dcd/owt/model.py

@@ -0,0 +1,601 @@
+import math
+import typing
+
+import einops
+import flash_attn
+import flash_attn.layers.rotary
+import huggingface_hub
+import omegaconf
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import transformers
+
+from .config import DUOConfig
+
+# Flags required to enable jit fusion kernels
+torch._C._jit_set_profiling_mode(False)
+torch._C._jit_set_profiling_executor(False)
+torch._C._jit_override_can_fuse_on_cpu(True)
+torch._C._jit_override_can_fuse_on_gpu(True)
+
+
+def bias_dropout_add_scale(
+    x: torch.Tensor,
+    bias: typing.Optional[torch.Tensor],
+    scale: torch.Tensor,
+    residual: typing.Optional[torch.Tensor],
+    prob: float,
+    training: bool) -> torch.Tensor:
+  if bias is not None:
+    out = scale * F.dropout(x + bias, p=prob, training=training)
+  else:
+    out = scale * F.dropout(x, p=prob, training=training)
+
+  if residual is not None:
+    out = residual + out
+  return out
+
+
+def get_bias_dropout_add_scale(training):
+  def _bias_dropout_add(x, bias, scale, residual, prob):
+    return bias_dropout_add_scale(
+      x, bias, scale, residual, prob, training)
+
+  return _bias_dropout_add
+
+
+# function overload
+def modulate(x: torch.Tensor,
+             shift: torch.Tensor,
+             scale: torch.Tensor) -> torch.Tensor:
+  return x * (1 + scale) + shift
+
+
+@torch.jit.script
+def bias_dropout_add_scale_fused_train(
+    x: torch.Tensor,
+    bias: typing.Optional[torch.Tensor],
+    scale: torch.Tensor,
+    residual: typing.Optional[torch.Tensor],
+    prob: float) -> torch.Tensor:
+  return bias_dropout_add_scale(
+    x, bias, scale, residual, prob, True)
+
+
+@torch.jit.script
+def bias_dropout_add_scale_fused_inference(
+    x: torch.Tensor,
+    bias: typing.Optional[torch.Tensor],
+    scale: torch.Tensor,
+    residual: typing.Optional[torch.Tensor],
+    prob: float) -> torch.Tensor:
+  return bias_dropout_add_scale(
+    x, bias, scale, residual, prob, False)
+
+
+@torch.jit.script
+def modulate_fused(x: torch.Tensor,
+                   shift: torch.Tensor,
+                   scale: torch.Tensor) -> torch.Tensor:
+  return modulate(x, shift, scale)
+
+
+class Rotary(torch.nn.Module):
+  def __init__(self, dim, base=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, seq_dim=1):
+    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.)
+      self.sin_cached[:,:,2,:,:].fill_(0.)
+
+    return self.cos_cached, self.sin_cached
+
+
+def rotate_half(x):
+  x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :]
+  return torch.cat((-x2, x1), dim=-1)
+
+
+def split_and_apply_rotary_pos_emb(qkv, rotary_cos_sin):
+  with torch.cuda.amp.autocast(enabled=False):
+    cos, sin = rotary_cos_sin
+    cos = cos.to(qkv.dtype)
+    sin = sin.to(qkv.dtype)
+    cos = cos[0,:,0,0,:cos.shape[-1]//2]
+    sin = sin[0,:,0,0,:sin.shape[-1]//2]
+    q, k, v = qkv.chunk(3, dim=2)
+    q = flash_attn.layers.rotary.apply_rotary_emb_torch(
+      q.squeeze(dim=2), cos, sin)
+    k = flash_attn.layers.rotary.apply_rotary_emb_torch(
+      k.squeeze(dim=2), cos, sin)
+    v = v.squeeze(dim=2)
+  return q, k, v
+
+
+def apply_rotary_pos_emb(qkv, cos, sin):
+  cos = cos[0,:,0,0,:cos.shape[-1]//2]
+  sin = sin[0,:,0,0,:sin.shape[-1]//2]
+  return flash_attn.layers.rotary.apply_rotary_emb_qkv_(qkv, cos, sin)
+
+
+def regular_attention_multi_headed(q, k, v):
+  # Assuming qkv is a tensor with shape [batch, seq_len, 3, num_heads, head_dim]
+  # where the 3 represents Q, K, V packed in that order
+  attention_output = F.scaled_dot_product_attention(
+    query=q.transpose(1, 2),
+    key=k.transpose(1, 2),
+    value=v.transpose(1, 2),
+    attn_mask=None,
+    dropout_p=0.0,
+    is_causal=False)
+  # [batch_size, seq_len, num_heads, head_dim]
+  attention_output = attention_output.transpose(1, 2)
+  return einops.rearrange(attention_output, 'b s h d -> b s (h d)')
+
+
+#################################################################################
+#                                  Layers                                       #
+#################################################################################
+class LayerNorm(nn.Module):
+  def __init__(self, dim):
+    super().__init__()
+    self.weight = nn.Parameter(torch.ones([dim]))
+    self.dim = dim
+  def forward(self, x):
+    with torch.cuda.amp.autocast(enabled=False):
+      x = F.layer_norm(x.float(), [self.dim])
+    return x * self.weight[None, None, :]
+
+
+def residual_linear(x, W, x_skip, residual_scale):
+  """x_skip + residual_scale * W @ x"""
+  dim_out, dim_in = W.shape[0], W.shape[1]
+  return torch.addmm(
+    x_skip.view(-1, dim_out),
+    x.view(-1, dim_in),
+    W.T,
+    alpha=residual_scale).view(*x.shape[:-1], dim_out)
+
+
+#################################################################################
+#               Embedding Layers for Timesteps and Class Labels                 #
+#################################################################################
+class TimestepEmbedder(nn.Module):
+  """
+  Embeds scalar timesteps into vector representations.
+  """
+  def __init__(self, hidden_size, frequency_embedding_size=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(t, dim, max_period=10000):
+    """
+    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.
+    """
+    # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
+    half = dim // 2
+    freqs = torch.exp(
+      - math.log(max_period)
+      * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device)
+      / half)
+    args = t[:, 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, t):
+    t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+    t_emb = self.mlp(t_freq)
+    return t_emb
+
+
+class LabelEmbedder(nn.Module):
+  """Embeds class labels into vector representations.
+  
+  Also handles label dropout for classifier-free guidance.
+  """
+  def __init__(self, num_classes, cond_size):
+    super().__init__()
+    self.embedding_table = nn.Embedding(num_classes + 1, cond_size)
+    self.num_classes = num_classes
+
+    # TODO think of initializing with 0.02 std deviation like in original DiT paper
+
+  def forward(self, labels):
+    embeddings = self.embedding_table(labels)
+    return embeddings
+    
+
+#################################################################################
+#                                 Core Model                                    #
+#################################################################################
+
+class DDiTBlockCausal(nn.Module):
+  def __init__(self, dim, n_heads, mlp_ratio=4, dropout=0.1):
+    super().__init__()
+    self.n_heads = n_heads
+
+    self.norm1 = LayerNorm(dim)
+    self.attn_qkv = nn.Linear(dim, 3 * dim, bias=False)
+    self.attn_out = nn.Linear(dim, dim, bias=False)
+    self.dropout1 = nn.Dropout(dropout)
+
+    self.norm2 = LayerNorm(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.dropout2 = nn.Dropout(dropout)
+    self.dropout = dropout
+
+  def _get_bias_dropout_scale(self):
+    if self.training:
+      return bias_dropout_add_scale_fused_train
+    else:
+      return bias_dropout_add_scale_fused_inference
+
+  def forward(self, x, rotary_cos_sin, **kwargs):
+    del kwargs
+    batch_size, seq_len = x.shape[0], x.shape[1]
+
+    bias_dropout_scale_fn = self._get_bias_dropout_scale()
+
+    # attention operation
+    x_skip = x
+    x = self.norm1(x)
+
+    qkv = self.attn_qkv(x)
+    qkv = einops.rearrange(
+      qkv,
+      'b s (three h d) -> b s three h d',
+      three=3,
+      h=self.n_heads)
+    with torch.cuda.amp.autocast(enabled=False):
+      cos, sin = rotary_cos_sin
+      qkv = apply_rotary_pos_emb(
+        qkv, cos.to(qkv.dtype), sin.to(qkv.dtype)
+      )
+    qkv = einops.rearrange(qkv, 'b s ... -> (b s) ...')
+    cu_seqlens = torch.arange(
+      0, (batch_size + 1) * seq_len,
+      step=seq_len, dtype=torch.int32, device=qkv.device)
+    x = flash_attn.flash_attn_interface.flash_attn_varlen_qkvpacked_func(
+      qkv, cu_seqlens, seq_len, 0.0, causal=True)
+
+    x = einops.rearrange(x, '(b s) h d -> b s (h d)',
+                         b=batch_size)
+
+    scale = torch.ones(1, device=x.device, dtype=x.dtype)
+    x = bias_dropout_scale_fn(
+      self.attn_out(x), None, scale, x_skip, self.dropout)
+
+    # mlp operation
+    x = bias_dropout_scale_fn(
+      self.mlp(self.norm2(x)), None, scale, x, self.dropout)
+    return x
+
+
+
+class DDiTBlock(nn.Module):
+  def __init__(self, dim, n_heads, adaLN,
+               cond_dim=None, mlp_ratio=4,
+               dropout=0.1):
+    super().__init__()
+    self.n_heads = n_heads
+    self.adaLN = adaLN
+
+    self.norm1 = LayerNorm(dim)
+    self.attn_qkv = nn.Linear(dim, 3 * dim, bias=False)
+    self.attn_out = nn.Linear(dim, dim, bias=False)
+    self.dropout1 = nn.Dropout(dropout)
+
+    self.norm2 = LayerNorm(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.dropout2 = nn.Dropout(dropout)
+    self.dropout = dropout
+
+    if self.adaLN:
+      self.adaLN_modulation = nn.Linear(cond_dim, 6 * dim)
+      self.adaLN_modulation.weight.data.zero_()
+      self.adaLN_modulation.bias.data.zero_()
+
+
+  def _get_bias_dropout_scale(self):
+    if self.training:
+      return bias_dropout_add_scale_fused_train
+    else:
+      return bias_dropout_add_scale_fused_inference
+
+
+  def forward(self, x, rotary_cos_sin, c=None):
+
+    bias_dropout_scale_fn = self._get_bias_dropout_scale()
+
+    x_skip = x
+    x = self.norm1(x)
+
+    if self.adaLN:
+      # self.adaLN_modulation(c): (128, 1536)
+      # self.adaLN_modulation(c)[:, None]: (128, 1, 1536)
+      # "" .chunk(6, dim=2) returns 6 tuples of shapes (128, 1, 256)
+      (shift_msa, scale_msa, gate_msa, shift_mlp,
+       scale_mlp, gate_mlp) = self.adaLN_modulation(c)[:, None].chunk(6, dim=2)
+      x = modulate_fused(x, shift_msa, scale_msa)
+
+    qkv = einops.rearrange(
+      self.attn_qkv(x),
+      'b s (three h d) -> b s three h d',
+      three=3,
+      h=self.n_heads)
+    q, k, v = split_and_apply_rotary_pos_emb(qkv, rotary_cos_sin)
+    
+    x = regular_attention_multi_headed(q, k, v)
+
+    if self.adaLN:
+      x = bias_dropout_scale_fn(self.attn_out(x),
+                                None,
+                                gate_msa,
+                                x_skip,
+                                self.dropout)
+      x = bias_dropout_scale_fn(
+        self.mlp(modulate_fused(
+          self.norm2(x), shift_mlp, scale_mlp)),
+        None, gate_mlp, x, self.dropout)
+    else:
+      scale = torch.ones(1, device=x.device, dtype=x.dtype)
+      x = bias_dropout_scale_fn(
+        self.attn_out(x), None, scale, x_skip, self.dropout)
+      x = bias_dropout_scale_fn(
+        self.mlp(self.norm2(x)), None, scale, x, self.dropout)
+    return x
+
+
+class EmbeddingLayer(nn.Module):
+  def __init__(self, dim, vocab_dim):
+    super().__init__()
+    self.embedding = nn.Parameter(torch.empty((vocab_dim, dim)))
+    torch.nn.init.kaiming_uniform_(self.embedding, a=math.sqrt(5))
+
+  def forward(self, x):
+    if x.ndim == 2:
+      return self.embedding[x]
+    assert x.ndim == 3
+    return torch.einsum(
+      "blv,ve->ble",
+      torch.nn.functional.softmax(x, dim=-1).float(),
+      self.embedding.float()).to(x.dtype)
+
+
+class DDiTFinalLayer(nn.Module):
+  def __init__(self, hidden_size, out_channels, cond_dim,
+               adaLN):
+    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 = adaLN
+    if self.adaLN:
+      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, c):
+    x = self.norm_final(x)
+    if self.adaLN:
+      shift, scale = self.adaLN_modulation(c)[:, None].chunk(2, dim=2)
+      x = modulate_fused(x, shift, scale)
+    x = self.linear(x)
+    return x
+
+
+class DIT(nn.Module, huggingface_hub.PyTorchModelHubMixin):
+  def __init__(self, config, vocab_size: int):
+    super().__init__()
+    if type(config) == dict:
+      config = omegaconf.OmegaConf.create(config)
+    self.causal = config.algo.causal_attention
+    self.adaLN = not self.causal
+    self.config = config
+    self.vocab_size = vocab_size
+    dim = config.model.hidden_size
+    cond_dim = config.model.cond_dim
+    self.vocab_embed = EmbeddingLayer(dim, vocab_size)
+    if not self.causal:
+      self.sigma_map = TimestepEmbedder(cond_dim)
+    self.rotary_emb = Rotary(dim // config.model.n_heads)
+
+    blocks = []
+    for _ in range(config.model.n_blocks):
+      if self.causal:
+        block = DDiTBlockCausal(
+          dim=dim,
+          n_heads=config.model.n_heads,
+          dropout=config.model.dropout)
+      else:
+        block = DDiTBlock(
+          dim=dim,
+          n_heads=config.model.n_heads,
+          cond_dim=cond_dim,
+          adaLN=self.adaLN,
+          dropout=config.model.dropout)
+      blocks.append(block)
+    self.blocks = nn.ModuleList(blocks)
+
+    self.output_layer = DDiTFinalLayer(
+      hidden_size=dim,
+      out_channels=vocab_size,
+      cond_dim=cond_dim,
+      adaLN=self.adaLN)
+    self.scale_by_sigma = config.model.scale_by_sigma
+
+  def _get_bias_dropout_scale(self):
+    if self.training:
+      return bias_dropout_add_scale_fused_train
+    else:
+      return  bias_dropout_add_scale_fused_inference
+
+  def forward(self, x, sigma):
+    x = self.vocab_embed(x)
+    if self.causal:
+      t_cond = None
+    else:
+      t_cond = F.silu(self.sigma_map(sigma))
+
+    rotary_cos_sin = self.rotary_emb(x).to(x.device)
+
+    with torch.cuda.amp.autocast(dtype=torch.bfloat16):
+      for i in range(len(self.blocks)):
+        x = self.blocks[i](x, rotary_cos_sin, c=t_cond)
+      x = self.output_layer(x, c=t_cond)
+
+    return x
+
+
+
+class HFDIT(torch.nn.Module):
+  def __init__(self, config):
+    super().__init__()
+    self.causal = config.causal
+    self.adaLN = not self.causal
+    self.vocab_size = config.vocab_size
+    dim = config.hidden_dim
+    cond_dim = config.cond_dim
+    self.vocab_embed = EmbeddingLayer(dim, self.vocab_size)
+    if not self.causal:
+      self.sigma_map = TimestepEmbedder(cond_dim)
+    self.rotary_emb = Rotary(dim // config.n_heads)
+
+    blocks = []
+    for _ in range(config.n_blocks):
+      if self.causal:
+        block = DDiTBlockCausal(
+          dim=dim,
+          n_heads=config.n_heads,
+          dropout=config.dropout)
+      else:
+        block = DDiTBlock(
+          dim=dim,
+          n_heads=config.n_heads,
+          cond_dim=cond_dim,
+          adaLN=self.adaLN,
+          dropout=config.dropout)
+      blocks.append(block)
+    self.blocks = torch.nn.ModuleList(blocks)
+
+    self.output_layer = DDiTFinalLayer(
+      hidden_size=dim,
+      out_channels=self.vocab_size,
+      cond_dim=cond_dim,
+      adaLN=self.adaLN)
+
+  def _get_bias_dropout_scale(self):
+    if self.training:
+      return bias_dropout_add_scale_fused_train
+    else:
+      return  bias_dropout_add_scale_fused_inference
+
+  def forward(self, x, sigma, output_hidden_states=False):
+    all_hidden_states = []
+    x = self.vocab_embed(x)
+    if output_hidden_states:
+      all_hidden_states.append(x)
+    if self.causal:
+      t_cond = None
+    else:
+      t_cond = F.silu(self.sigma_map(sigma))
+
+    rotary_cos_sin = self.rotary_emb(x)
+    with torch.cuda.amp.autocast(
+      dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float32):
+      for i in range(len(self.blocks)):
+        x = self.blocks[i](x, rotary_cos_sin, c=t_cond)
+        if output_hidden_states:
+          all_hidden_states.append(x)
+      x = self.output_layer(x, c=t_cond)
+    return x, all_hidden_states
+
+
+class DUO(transformers.PreTrainedModel):
+  """HF-compatible model."""
+  config_class = DUOConfig
+  base_model_prefix = 'duo'
+
+  def __init__(self, config: DUOConfig):
+    super().__init__(config)
+    self.config = config
+    self.backbone = HFDIT(config)
+
+  def reset_kv_cache(self):
+    for block in self.backbone.blocks:
+      block.kv_cache = None
+
+  def forward(
+      self,
+      input_ids: torch.LongTensor = None,
+      timesteps: torch.FloatTensor = None,
+      output_hidden_states: typing.Optional[bool] = None,
+      return_dict: typing.Optional[bool] = None,
+  ) -> typing.Union[
+    torch.Tensor, typing.Tuple,
+    transformers.modeling_outputs.MaskedLMOutput]:
+    """HF-compatible forward method."""
+    output_hidden_states = (
+      output_hidden_states
+      if output_hidden_states is not None
+      else self.config.output_hidden_states
+    )
+    return_dict = return_dict \
+      if return_dict is not None \
+      else self.config.use_return_dict
+    
+    logits, all_hidden_states = self.backbone(
+      x=input_ids,
+      sigma=timesteps,
+      output_hidden_states=output_hidden_states,
+    )
+    if return_dict:
+      return transformers.modeling_outputs.MaskedLMOutput(
+        logits=logits,
+        hidden_states=all_hidden_states if output_hidden_states else None,
+        loss=None
+      )
+    elif output_hidden_states:
+      return logits, all_hidden_states
+    else:
+      return logits

checkpoints/quarantine/20260818T074907Z/official/mdlm/owt/config.json

@@ -0,0 +1,22 @@
+{
+  "_name_or_path": "kuleshov-group/mdlm-owt",
+  "architectures": [
+    "MDLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_mdlm.MDLMConfig",
+    "AutoModelForMaskedLM": "modeling_mdlm.MDLM"
+  },
+  "cond_dim": 128,
+  "dropout": 0.1,
+  "hidden_dim": 768,
+  "model_length": 1024,
+  "model_type": "mdlm",
+  "n_blocks": 12,
+  "n_heads": 12,
+  "return_dict": false,
+  "time_conditioning": false,
+  "torch_dtype": "float32",
+  "transformers_version": "4.38.2",
+  "vocab_size": 50258
+}

checkpoints/quarantine/20260818T074907Z/official/mdlm/owt/configuration_mdlm.py

@@ -0,0 +1,31 @@
+"""MDLM config for Hugging Face.
+
+"""
+
+import transformers
+
+
+class MDLMConfig(transformers.PretrainedConfig):
+  """Hugging Face configuration class for MDLM."""
+  model_type = "mdlm"
+
+  def __init__(
+    self,
+    vocab_size: int = 50258,
+    model_length: int = 1024,
+    hidden_dim: int = 768,
+    cond_dim: int = 129,
+    n_blocks: int = 12,
+    n_heads: int = 12,
+    dropout: float = 0.1,
+    time_conditioning: bool = False,
+    ** kwargs):
+    super().__init__(**kwargs)
+    self.vocab_size = vocab_size
+    self.model_length = model_length
+    self.hidden_dim = hidden_dim
+    self.cond_dim = cond_dim
+    self.n_blocks = n_blocks
+    self.n_heads = n_heads
+    self.dropout = dropout
+    self.time_conditioning = time_conditioning

checkpoints/quarantine/20260820T072002Z/official/flm/owt/README.md

@@ -0,0 +1,45 @@
+---
+tags:
+- model_hub_mixin
+- pytorch_model_hub_mixin
+license: apache-2.0
+datasets:
+- Skylion007/openwebtext
+language:
+- en
+---
+  
+  ## FLM (OpenWebText)
+
+  For more information and implementation details, visit our github repository: [FLM](https://github.com/david3684/flm).
+
+  ## Model Details
+
+  The model has a context length of 1024, trained for 1M steps on the OpenWebText dataset.
+
+  For more details, please see our paper: 
+  - [arXiv:2602.16813](https://arxiv.org/abs/2602.16813)
+  - [Huggingface Papers](https://huggingface.co/papers/2602.16813)
+
+  ## Citation
+
+  Please cite our work using the bibtex below:
+
+  **BibTeX:**
+
+
+  ```bibtex
+  @article{lee2026flow,
+      title={Flow Map Language Models: One-step Language Modeling via Continuous Denoising},
+      author={Chanhyuk Lee and Jaehoon Yoo and Manan Agarwal
+              and Sheel Shah and Jerry Huang
+              and Aditi Raghunathan and Seunghoon Hong
+              and Nicholas M. Boffi and Jinwoo Kim},
+      journal={arXiv preprint arXiv:2602.16813},
+      year={2026},
+  }
+  ```
+
+  ## Model Card Contact
+
+  Chanhyuk Lee (chanhyuk-lee@kaist.ac.kr)

checkpoints/quarantine/20260820T072143Z/official/langflow/lm1b/config.json

@@ -0,0 +1,29 @@
+{
+  "_name_or_path": "nealchen/langflow-lm1b",
+  "architectures": [
+    "LangFlow"
+  ],
+  "auto_map": {
+    "AutoConfig": "config.LangFlowConfig",
+    "AutoModelForMaskedLM": "model.LangFlow"
+  },
+  "model_type": "LangFlow",
+  "vocab_size": 30522,
+  "hidden_size": 768,
+  "cond_dim": 128,
+  "n_blocks": 12,
+  "n_heads": 12,
+  "dropout": 0.1,
+  "model_length": 128,
+  "use_normalized_embedding": true,
+  "embedding_norm_method": "layernorm",
+  "self_conditioning": true,
+  "use_bias": true,
+  "gumbel_loc": 4.723,
+  "gumbel_scale": 0.852,
+  "gumbel_cutoff": 1e-5,
+  "gumbel_entropy": 7.02,
+  "return_dict": true,
+  "torch_dtype": "float32",
+  "transformers_version": "4.38.2"
+}

checkpoints/quarantine/20260820T072143Z/official/langflow/lm1b/model.py

@@ -0,0 +1,560 @@
+"""HuggingFace model implementation for LangFlow.
+
+LangFlow is a continuous diffusion language model that operates in embedding space.
+"""
+
+import math
+import typing
+
+import einops
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import transformers
+
+from .config import LangFlowConfig
+
+
+# Flags required to enable jit fusion kernels
+torch._C._jit_set_profiling_mode(False)
+torch._C._jit_set_profiling_executor(False)
+torch._C._jit_override_can_fuse_on_cpu(True)
+torch._C._jit_override_can_fuse_on_gpu(True)
+
+
+def bias_dropout_add_scale(
+    x: torch.Tensor,
+    bias: typing.Optional[torch.Tensor],
+    scale: torch.Tensor,
+    residual: typing.Optional[torch.Tensor],
+    prob: float,
+    training: bool) -> torch.Tensor:
+    if bias is not None:
+        out = scale * F.dropout(x + bias, p=prob, training=training)
+    else:
+        out = scale * F.dropout(x, p=prob, training=training)
+
+    if residual is not None:
+        out = residual + out
+    return out
+
+
+@torch.jit.script
+def bias_dropout_add_scale_fused_train(
+    x: torch.Tensor,
+    bias: typing.Optional[torch.Tensor],
+    scale: torch.Tensor,
+    residual: typing.Optional[torch.Tensor],
+    prob: float) -> torch.Tensor:
+    return bias_dropout_add_scale(x, bias, scale, residual, prob, True)
+
+
+@torch.jit.script
+def bias_dropout_add_scale_fused_inference(
+    x: torch.Tensor,
+    bias: typing.Optional[torch.Tensor],
+    scale: torch.Tensor,
+    residual: typing.Optional[torch.Tensor],
+    prob: float) -> torch.Tensor:
+    return bias_dropout_add_scale(x, bias, scale, residual, prob, False)
+
+
+@torch.jit.script
+def modulate_fused(x: torch.Tensor,
+                   shift: torch.Tensor,
+                   scale: torch.Tensor) -> torch.Tensor:
+    return x * (1 + scale) + shift
+
+
+class Rotary(nn.Module):
+    def __init__(self, dim, base=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, seq_dim=1):
+        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)
+            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)
+            self.cos_cached[:, :, 2, :, :].fill_(1.)
+            self.sin_cached[:, :, 2, :, :].fill_(0.)
+        return self.cos_cached, self.sin_cached
+
+
+def _apply_rotary_emb(x, cos, sin):
+    # x: [batch, seqlen, nheads, headdim]
+    # cos, sin: [seqlen, headdim//2]
+    ro_dim = cos.shape[-1] * 2
+    # Expand to [1, seqlen, 1, ro_dim] for broadcasting
+    cos = torch.cat([cos, cos], dim=-1)[None, :, None, :]
+    sin = torch.cat([sin, sin], dim=-1)[None, :, None, :]
+    x_rot = x[..., :ro_dim]
+    x1, x2 = x_rot.chunk(2, dim=-1)
+    x_rotated = torch.cat([-x2, x1], dim=-1)
+    return torch.cat([x_rot * cos + x_rotated * sin, x[..., ro_dim:]], dim=-1)
+
+
+def split_and_apply_rotary_pos_emb(qkv, rotary_cos_sin):
+    with torch.autocast(device_type='cuda', enabled=False):
+        cos, sin = rotary_cos_sin
+        cos = cos.to(qkv.dtype)
+        sin = sin.to(qkv.dtype)
+        cos = cos[0, :, 0, 0, :cos.shape[-1]//2]
+        sin = sin[0, :, 0, 0, :sin.shape[-1]//2]
+        q, k, v = qkv.chunk(3, dim=2)
+        q = _apply_rotary_emb(q.squeeze(dim=2), cos, sin)
+        k = _apply_rotary_emb(k.squeeze(dim=2), cos, sin)
+        v = v.squeeze(dim=2)
+    return q, k, v
+
+
+def regular_attention_multi_headed(q, k, v):
+    attention_output = F.scaled_dot_product_attention(
+        query=q.transpose(1, 2),
+        key=k.transpose(1, 2),
+        value=v.transpose(1, 2),
+        attn_mask=None,
+        dropout_p=0.0,
+        is_causal=False)
+    attention_output = attention_output.transpose(1, 2)
+    return einops.rearrange(attention_output, 'b s h d -> b s (h d)')
+
+
+class LayerNorm(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones([dim]))
+        self.dim = dim
+
+    def forward(self, x):
+        with torch.autocast(device_type='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, frequency_embedding_size=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(t, dim, max_period=10000):
+        half = dim // 2
+        freqs = torch.exp(
+            -math.log(max_period)
+            * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device)
+            / half)
+        args = t[:, 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, t):
+        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+        t_emb = self.mlp(t_freq)
+        return t_emb
+
+
+class DDiTBlock(nn.Module):
+    def __init__(self, dim, n_heads, cond_dim, mlp_ratio=4, dropout=0.1):
+        super().__init__()
+        self.n_heads = n_heads
+
+        self.norm1 = LayerNorm(dim)
+        self.attn_qkv = nn.Linear(dim, 3 * dim, bias=False)
+        self.attn_out = nn.Linear(dim, dim, bias=False)
+
+        self.norm2 = LayerNorm(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.dropout = dropout
+
+        self.adaLN_modulation = nn.Linear(cond_dim, 6 * dim)
+        self.adaLN_modulation.weight.data.zero_()
+        self.adaLN_modulation.bias.data.zero_()
+
+    def _get_bias_dropout_scale(self):
+        if self.training:
+            return bias_dropout_add_scale_fused_train
+        else:
+            return bias_dropout_add_scale_fused_inference
+
+    def forward(self, x, rotary_cos_sin, c):
+        bias_dropout_scale_fn = self._get_bias_dropout_scale()
+
+        x_skip = x
+        x = self.norm1(x)
+
+        (shift_msa, scale_msa, gate_msa, shift_mlp,
+         scale_mlp, gate_mlp) = self.adaLN_modulation(c)[:, None].chunk(6, dim=2)
+        x = modulate_fused(x, shift_msa, scale_msa)
+
+        qkv = einops.rearrange(
+            self.attn_qkv(x),
+            'b s (three h d) -> b s three h d',
+            three=3,
+            h=self.n_heads)
+        q, k, v = split_and_apply_rotary_pos_emb(qkv, rotary_cos_sin)
+        x = regular_attention_multi_headed(q, k, v)
+
+        x = bias_dropout_scale_fn(self.attn_out(x), None, gate_msa, x_skip, self.dropout)
+        x = bias_dropout_scale_fn(
+            self.mlp(modulate_fused(self.norm2(x), shift_mlp, scale_mlp)),
+            None, gate_mlp, x, self.dropout)
+        return x
+
+
+def _normalize_embedding_layernorm(weight: torch.Tensor) -> torch.Tensor:
+    """Normalize embedding weights to unit norm per row, then scale by sqrt(dim)."""
+    normalized = F.normalize(weight.float(), dim=-1)
+    return (normalized * math.sqrt(weight.shape[-1])).to(weight.dtype)
+
+
+class EmbeddingLayer(nn.Module):
+    """Embedding layer with optional layernorm normalization."""
+    
+    def __init__(self, dim, vocab_dim, use_normalized_embedding=True):
+        super().__init__()
+        self.dim = dim
+        self.vocab_dim = vocab_dim
+        self.use_normalized_embedding = use_normalized_embedding
+        self.embedding = nn.Parameter(torch.empty((vocab_dim, dim)))
+        nn.init.kaiming_uniform_(self.embedding, a=math.sqrt(5))
+
+    def _get_embedding(self):
+        if self.use_normalized_embedding:
+            return _normalize_embedding_layernorm(self.embedding)
+        return self.embedding
+
+    def forward(self, x):
+        embedding = self._get_embedding()
+        if x.ndim == 2:
+            return embedding[x]
+        assert x.ndim == 3  # probabilities
+        return torch.einsum("blv,ve->ble", x.float(), embedding.float()).to(x.dtype)
+
+
+class DDiTFinalLayer(nn.Module):
+    def __init__(self, hidden_size, out_channels, cond_dim):
+        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, c):
+        x = self.norm_final(x)
+        shift, scale = self.adaLN_modulation(c)[:, None].chunk(2, dim=2)
+        x = modulate_fused(x, shift, scale)
+        x = self.linear(x)
+        return x
+
+
+class GumbelProposal(nn.Module):
+    """Learnable Gumbel distribution proposal for sampling gamma (log-SNR)."""
+    
+    def __init__(self, loc: float = 4.723, scale: float = 0.852, 
+                 cutoff: float = 1e-5, entropy: float = 7.02):
+        super().__init__()
+        self.loc = nn.Parameter(torch.tensor(loc))
+        self.scale = nn.Parameter(torch.tensor(scale))
+        self.cutoff = cutoff
+        self.entropy = nn.Parameter(torch.tensor(entropy))
+
+    def _get_distribution(self) -> torch.distributions.Gumbel:
+        return torch.distributions.Gumbel(self.loc, self.scale)
+
+    @property
+    def gamma_min(self) -> float:
+        return float(self.loc - math.log(-math.log(self.cutoff)) * self.scale)
+
+    @property
+    def gamma_max(self) -> float:
+        return float(self.loc - math.log(self.cutoff) * self.scale)
+
+    def forward(self, q: torch.Tensor) -> torch.Tensor:
+        """Convert uniform samples to gamma values via inverse CDF."""
+        gamma = self._get_distribution().icdf(q)
+        return gamma.clamp(min=self.gamma_min, max=self.gamma_max)
+
+    def log_pdf(self, gamma: torch.Tensor) -> torch.Tensor:
+        """Compute log probability density at gamma."""
+        return self._get_distribution().log_prob(gamma)
+
+
+class LangFlowBackbone(nn.Module):
+    """DiT backbone for LangFlow."""
+    
+    def __init__(self, config: LangFlowConfig):
+        super().__init__()
+        self.config = config
+        dim = config.hidden_size
+        cond_dim = config.cond_dim
+
+        self.vocab_embed = EmbeddingLayer(
+            dim, config.vocab_size,
+            use_normalized_embedding=config.use_normalized_embedding)
+        self.sigma_map = TimestepEmbedder(cond_dim)
+        self.rotary_emb = Rotary(dim // config.n_heads)
+
+        self.blocks = nn.ModuleList([
+            DDiTBlock(dim=dim, n_heads=config.n_heads, cond_dim=cond_dim, dropout=config.dropout)
+            for _ in range(config.n_blocks)
+        ])
+
+        self.output_layer = DDiTFinalLayer(
+            hidden_size=dim, out_channels=config.vocab_size, cond_dim=cond_dim)
+
+        # Self-conditioning projection
+        if config.self_conditioning:
+            self.self_cond_proj = nn.Linear(dim * 2, dim, bias=False)
+            nn.init.zeros_(self.self_cond_proj.weight)
+
+    def forward(self, x_embed, sigma, x_self_cond=None, output_hidden_states=False):
+        """Forward pass from embeddings.
+        
+        Args:
+            x_embed: [B, L, D] - Input embeddings (possibly noisy)
+            sigma: [B] - Gamma values (log-SNR)
+            x_self_cond: [B, L, D] - Self-conditioning embeddings (optional)
+            output_hidden_states: Whether to return all hidden states
+        
+        Returns:
+            logits: [B, L, vocab_size]
+            hidden_states: List of hidden states if output_hidden_states=True
+        """
+        all_hidden_states = []
+        x = x_embed
+        
+        if output_hidden_states:
+            all_hidden_states.append(x)
+
+        # Self-conditioning
+        if self.config.self_conditioning:
+            if x_self_cond is None:
+                x_self_cond = torch.zeros_like(x)
+            x = x + self.self_cond_proj(torch.cat([x, x_self_cond], dim=-1))
+
+        t_cond = F.silu(self.sigma_map(sigma))
+        rotary_cos_sin = self.rotary_emb(x)
+
+        with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
+            for block in self.blocks:
+                x = block(x, rotary_cos_sin, c=t_cond)
+                if output_hidden_states:
+                    all_hidden_states.append(x)
+            x = self.output_layer(x, c=t_cond)
+
+        return x, all_hidden_states
+
+
+class LangFlow(transformers.PreTrainedModel):
+    """HuggingFace-compatible LangFlow model.
+    
+    LangFlow is a continuous diffusion language model that operates in embedding space.
+    It uses a DiT (Diffusion Transformer) backbone with:
+    - Self-conditioning: uses previous predictions as additional input
+    - Bias (preconditioning): skip connection for improved generation
+    - Normalized embeddings: layernorm on embedding vectors
+    - Learnable Gumbel proposal for gamma (log-SNR) sampling
+    """
+    config_class = LangFlowConfig
+    base_model_prefix = "langflow"
+
+    def __init__(self, config: LangFlowConfig):
+        super().__init__(config)
+        self.config = config
+        self.backbone = LangFlowBackbone(config)
+        self.proposal = GumbelProposal(
+            loc=config.gumbel_loc,
+            scale=config.gumbel_scale,
+            cutoff=config.gumbel_cutoff,
+            entropy=config.gumbel_entropy)
+
+    def _get_embedding_matrix(self) -> torch.Tensor:
+        """Get the embedding matrix for bias skip connection."""
+        return self.backbone.vocab_embed._get_embedding()
+
+    def _embed_tokens(self, x: torch.Tensor) -> torch.Tensor:
+        """Embed tokens or probabilities to continuous embeddings."""
+        return self.backbone.vocab_embed(x)
+
+    def _forward_diffusion(self, x_embed: torch.Tensor, 
+                           gamma: torch.Tensor) -> torch.Tensor:
+        """Add noise to embeddings (forward diffusion process)."""
+        gamma = gamma.float()
+        alpha = torch.sigmoid(-gamma).sqrt()[:, None, None]
+        sigma = torch.sigmoid(gamma).sqrt()[:, None, None]
+        noise = torch.randn_like(x_embed)
+        return (x_embed * alpha + noise * sigma).to(x_embed.dtype)
+
+    def _euler_edm_step(self, z: torch.Tensor, x_pred: torch.Tensor,
+                        t: torch.Tensor, s: torch.Tensor) -> torch.Tensor:
+        """Single Euler step for EDM sampling."""
+        t_ = t.double()
+        s_ = s.double()
+        cur = z.double() * ((F.softplus(t_) - F.softplus(s_)) / 2).exp()
+        end = torch.sigmoid(-s_).sqrt() * x_pred.double()
+        z = end.lerp(cur, ((s_ - t_) / 2).exp()).to(z.dtype)
+        return z
+
+    def forward(
+        self,
+        input_ids: typing.Optional[torch.LongTensor] = None,
+        noisy_embeds: typing.Optional[torch.FloatTensor] = None,
+        timesteps: typing.Optional[torch.FloatTensor] = None,
+        x_self_cond: typing.Optional[torch.FloatTensor] = None,
+        output_hidden_states: typing.Optional[bool] = None,
+        return_dict: typing.Optional[bool] = None,
+    ) -> typing.Union[torch.Tensor, typing.Tuple, transformers.modeling_outputs.MaskedLMOutput]:
+        """Forward pass for LangFlow.
+        
+        Args:
+            input_ids: [B, L] - Token IDs (will be embedded and noised if timesteps provided)
+            noisy_embeds: [B, L, D] - Pre-noised embeddings (alternative to input_ids)
+            timesteps: [B] - Gamma values (log-SNR) for conditioning
+            x_self_cond: [B, L, D] - Self-conditioning embeddings
+            output_hidden_states: Whether to return hidden states
+            return_dict: Whether to return MaskedLMOutput
+        
+        Returns:
+            logits or MaskedLMOutput
+        """
+        output_hidden_states = output_hidden_states if output_hidden_states is not None else False
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # Get embeddings
+        if noisy_embeds is not None:
+            z = noisy_embeds
+        elif input_ids is not None:
+            x_embed = self._embed_tokens(input_ids)
+            if timesteps is not None:
+                z = self._forward_diffusion(x_embed, timesteps)
+            else:
+                z = x_embed
+        else:
+            raise ValueError("Either input_ids or noisy_embeds must be provided")
+
+        if timesteps is None:
+            # Use minimum gamma for clean input
+            timesteps = torch.full((z.shape[0],), self.proposal.gamma_min, device=z.device)
+
+        # Process sigma
+        sigma = timesteps
+        if sigma.ndim == 2:
+            sigma = sigma.mean(-1)
+
+        # Get model output
+        logits, all_hidden_states = self.backbone(
+            z, sigma, x_self_cond=x_self_cond, output_hidden_states=output_hidden_states)
+
+        # Add bias (preconditioning) skip connection
+        if self.config.use_bias:
+            c_skip = ((F.softplus(-sigma) - sigma) / 2).exp()
+            embedding = self._get_embedding_matrix()
+            skip_logits = torch.matmul(z.float(), embedding.t().float())
+            logits = logits + c_skip[:, None, None] * skip_logits.to(logits.dtype)
+
+        if return_dict:
+            return transformers.modeling_outputs.MaskedLMOutput(
+                logits=logits,
+                hidden_states=all_hidden_states if output_hidden_states else None,
+                loss=None)
+        elif output_hidden_states:
+            return logits, all_hidden_states
+        else:
+            return logits
+
+    @torch.no_grad()
+    def generate_samples(
+        self,
+        num_samples: int = 1,
+        seq_length: typing.Optional[int] = None,
+        num_steps: int = 128,
+        device: typing.Optional[torch.device] = None,
+    ) -> torch.LongTensor:
+        """Generate samples using Euler-EDM solver.
+        
+        Args:
+            num_samples: Number of samples to generate
+            seq_length: Sequence length (defaults to config.model_length)
+            num_steps: Number of denoising steps
+            device: Device to generate on
+        
+        Returns:
+            samples: [num_samples, seq_length] - Generated token IDs
+        """
+        if seq_length is None:
+            seq_length = self.config.model_length
+        if device is None:
+            device = next(self.parameters()).device
+
+        embed_dim = self.config.hidden_size
+        eps = 1e-5
+
+        # Initialize with Gaussian noise
+        z = torch.randn(num_samples, seq_length, embed_dim, device=device)
+
+        # Create gamma schedule from t=1-eps to t=eps
+        t = torch.linspace(1.0 - eps, eps, num_steps, device=device)
+        gamma = self.proposal(t)
+
+        # Self-conditioning state
+        x_self_cond = None
+
+        # Euler-EDM sampling loop
+        for i in range(len(gamma) - 1):
+            gamma_t = gamma[i]
+            gamma_s = gamma[i + 1]
+
+            # Get model prediction
+            gamma_expanded = gamma_t.unsqueeze(0).expand(num_samples)
+            logits = self.forward(
+                noisy_embeds=z,
+                timesteps=gamma_expanded,
+                x_self_cond=x_self_cond,
+                return_dict=False)
+
+            # Convert logits to embedding prediction
+            probs = F.softmax(logits.float(), dim=-1)
+            x_pred = self._embed_tokens(probs)
+
+            # Update self-conditioning
+            if self.config.self_conditioning:
+                x_self_cond = x_pred
+
+            # Euler step
+            z = self._euler_edm_step(z, x_pred, gamma_t, gamma_s)
+
+        # Final step: get logits and take argmax
+        gamma_final = gamma[-1]
+        gamma_expanded = gamma_final.unsqueeze(0).expand(num_samples)
+        logits = self.forward(
+            noisy_embeds=z,
+            timesteps=gamma_expanded,
+            x_self_cond=x_self_cond,
+            return_dict=False)
+        samples = logits.argmax(dim=-1)
+
+        return samples

checkpoints/quarantine/20260821T024711Z/official/flm/lm1b/config.json

@@ -0,0 +1,209 @@
+{
+  "config": {
+    "algo": {
+      "T": 0,
+      "backbone": "dit",
+      "causal_attention": false,
+      "curriculum_end": 1000000,
+      "curriculum_start": 0,
+      "double_temb": false,
+      "flow_loss_type": "ce",
+      "flow_ratio": 0.75,
+      "gumbel_tau_log10_end": -0.0,
+      "gumbel_tau_log10_start": -0.0,
+      "ignore_bos": false,
+      "jvp_api": "autograd",
+      "loss_scale": "none",
+      "n_separated_blocks": -1,
+      "name": "dos",
+      "parameterization": "mean",
+      "pred_log_interval": 2000,
+      "sigma_min": 0.0001,
+      "subs_masking": false,
+      "t_max": 1.0,
+      "t_min": 0.0,
+      "time_condition": "alpha_t",
+      "time_conditioning": true,
+      "use_curriculum": true,
+      "use_discrete_schedule": true
+    },
+    "callbacks": {
+      "checkpoint_every_n_steps": {
+        "_target_": "lightning.pytorch.callbacks.ModelCheckpoint",
+        "auto_insert_metric_name": false,
+        "dirpath": "/home/david3684/projects/flm/hf_upload/checkpoints",
+        "every_n_train_steps": 50000,
+        "save_last": false,
+        "save_top_k": -1,
+        "verbose": true
+      },
+      "checkpoint_monitor": {
+        "_target_": "lightning.pytorch.callbacks.ModelCheckpoint",
+        "auto_insert_metric_name": false,
+        "dirpath": "/home/david3684/projects/flm/hf_upload/checkpoints",
+        "filename": "best_nll",
+        "mode": "min",
+        "monitor": "val/nll",
+        "save_last": true,
+        "save_top_k": 1,
+        "verbose": true
+      },
+      "learning_rate_monitor": {
+        "_target_": "lightning.pytorch.callbacks.LearningRateMonitor",
+        "logging_interval": "step"
+      }
+    },
+    "checkpointing": {
+      "resume_ckpt_path": "/home/work/RADAR/workspace/KAIST/discrete-mean-flow/text/outputs/lm1b/2026.01.05/154016/checkpoints/47-650000.ckpt",
+      "resume_from_ckpt": true,
+      "save_dir": "/home/david3684/projects/flm/hf_upload"
+    },
+    "data": {
+      "cache_dir": "/home/work/RADAR/workspace/KAIST/cdBDD/datasets/lm1b",
+      "insert_train_eos": true,
+      "insert_valid_eos": true,
+      "streaming": false,
+      "tokenizer_name_or_path": "bert-base-uncased",
+      "train": "lm1b",
+      "valid": "lm1b",
+      "wrap": true
+    },
+    "eval": {
+      "checkpoint_path": "",
+      "compute_generative_perplexity": true,
+      "compute_perplexity_on_sanity": false,
+      "disable_ema": false,
+      "gen_ppl_eval_model_name_or_path": "gpt2-large",
+      "generate_samples": true,
+      "generated_samples_path": "/home/david3684/projects/flm/hf_upload/samples.json",
+      "perplexity_batch_size": 8
+    },
+    "is_di4c": false,
+    "is_di4c_deterministic": true,
+    "loader": {
+      "batch_size": 128,
+      "eval_batch_size": 128,
+      "eval_global_batch_size": 512,
+      "global_batch_size": 512,
+      "num_workers": 192,
+      "pin_memory": true
+    },
+    "lr_scheduler": {
+      "_target_": "transformers.get_constant_schedule_with_warmup",
+      "num_warmup_steps": 2500
+    },
+    "mode": "train",
+    "model": {
+      "cond_dim": 128,
+      "dropout": 0.1,
+      "hidden_size": 768,
+      "length": 128,
+      "n_blocks": 12,
+      "n_heads": 12,
+      "name": "small",
+      "scale_by_sigma": true,
+      "tie_word_embeddings": false,
+      "type": "ddit",
+      "vocab_lookup": true
+    },
+    "noise": {
+      "denoiser_latent_conditioning": -1,
+      "eps": 0,
+      "freeze_decoder": false,
+      "freeze_encoder": false,
+      "parameterization": "log-linear",
+      "type": "log-linear"
+    },
+    "optim": {
+      "beta1": 0.9,
+      "beta2": 0.999,
+      "eps": 1e-08,
+      "ln_tune": "none",
+      "lr": 0.0003,
+      "use_muon": false,
+      "weight_decay": 0
+    },
+    "prior": {
+      "latent_height": 0,
+      "latent_width": 0,
+      "type": "none"
+    },
+    "sampling": {
+      "argmax_correction": false,
+      "hard_start": false,
+      "noise_removal": "meanflow",
+      "noise_removal_list": [
+        "uniform",
+        "uniform_alpha"
+      ],
+      "num_reflow_samples": 50000,
+      "num_sample_batches": 1,
+      "num_sample_log": 2,
+      "num_strides": 1,
+      "p_nucleus": 1.0,
+      "predictor": "ancestral",
+      "semi_ar": false,
+      "solver": "euler",
+      "steps": [
+        1024
+      ],
+      "stride_length": 1,
+      "tau_log10": 0.0,
+      "temperature": 1.0,
+      "use_float64": true
+    },
+    "seed": 1,
+    "strategy": {
+      "_target_": "lightning.pytorch.strategies.DDPStrategy",
+      "find_unused_parameters": true
+    },
+    "trainer": {
+      "_target_": "lightning.Trainer",
+      "accelerator": "cuda",
+      "accumulate_grad_batches": 1,
+      "devices": 8,
+      "gradient_clip_val": 1.0,
+      "limit_train_batches": 1.0,
+      "limit_val_batches": 1.0,
+      "log_every_n_steps": 100,
+      "max_steps": 1000000,
+      "num_nodes": 1,
+      "num_sanity_val_steps": 2,
+      "precision": "bf16",
+      "val_check_interval": 10000
+    },
+    "training": {
+      "antithetic_sampling": true,
+      "change_of_variables": false,
+      "double_ema": 0.0,
+      "ema": 0.9999,
+      "finetune_path": "",
+      "importance_sampling": false,
+      "load_ema": false,
+      "loss_precision": "bf16",
+      "loss_type": "flow",
+      "loss_weighting": false,
+      "not_sampling_t": false,
+      "pred_type": "x0",
+      "sampling_eps": 0.001,
+      "t_curriculum_max": 1.0,
+      "t_curriculum_min": 0.0,
+      "t_curriculum_steps": 100000,
+      "torch_compile_mode": "default",
+      "use_torch_compile": true
+    },
+    "wandb": {
+      "group": null,
+      "id": "fc867e70",
+      "job_type": null,
+      "name": "lm1b_full_flow_no_input_softmax_no_loss_scale_ce_loss_fc867e70",
+      "notes": "DOS",
+      "project": "lm1b_full",
+      "tags": [
+        "lm1b",
+        "dos"
+      ]
+    }
+  },
+  "vocab_size": 30522
+}

checkpoints/quarantine/20260821T024727Z/official/flm/owt/config.json

@@ -0,0 +1,208 @@
+{
+  "config": {
+    "algo": {
+      "T": 0,
+      "backbone": "dit",
+      "causal_attention": false,
+      "curriculum_end": 1000000,
+      "curriculum_start": 0,
+      "double_temb": false,
+      "flow_loss_type": "ce",
+      "flow_ratio": 0.75,
+      "gumbel_tau_log10_end": -0.0,
+      "gumbel_tau_log10_start": -0.0,
+      "ignore_bos": false,
+      "jvp_api": "autograd",
+      "learnable_loss_weighting": false,
+      "loss_scale": "none",
+      "n_separated_blocks": -1,
+      "name": "dos",
+      "parameterization": "mean",
+      "pred_log_interval": 2000,
+      "sigma_min": 0.0001,
+      "subs_masking": false,
+      "t_max": 1.0,
+      "t_min": 0.0,
+      "time_condition": "alpha_t",
+      "time_conditioning": true,
+      "use_curriculum": true,
+      "use_discrete_schedule": true
+    },
+    "callbacks": {
+      "checkpoint_every_n_steps": {
+        "_target_": "lightning.pytorch.callbacks.ModelCheckpoint",
+        "auto_insert_metric_name": false,
+        "dirpath": "/home/david3684/projects/flm/hf_upload/checkpoints",
+        "every_n_train_steps": 25000,
+        "save_last": false,
+        "save_top_k": -1,
+        "verbose": true
+      },
+      "checkpoint_monitor": {
+        "_target_": "lightning.pytorch.callbacks.ModelCheckpoint",
+        "auto_insert_metric_name": false,
+        "dirpath": "/home/david3684/projects/flm/hf_upload/checkpoints",
+        "filename": "best_nll",
+        "mode": "min",
+        "monitor": "val/nll",
+        "save_last": true,
+        "save_top_k": 1,
+        "verbose": true
+      },
+      "learning_rate_monitor": {
+        "_target_": "lightning.pytorch.callbacks.LearningRateMonitor",
+        "logging_interval": "step"
+      }
+    },
+    "checkpointing": {
+      "resume_ckpt_path": "/home/work/RADAR/workspace/KAIST/discrete-mean-flow/text/outputs/openwebtext-train/2026.02.05/005839/checkpoints/50-825000.ckpt",
+      "resume_from_ckpt": true,
+      "save_dir": "/home/david3684/projects/flm/hf_upload"
+    },
+    "data": {
+      "cache_dir": "/home/work/RADAR/workspace/KAIST/cdBDD/datasets/openwebtext",
+      "insert_train_eos": true,
+      "insert_valid_eos": true,
+      "streaming": false,
+      "tokenizer_name_or_path": "gpt2",
+      "train": "openwebtext-train",
+      "valid": "openwebtext-valid",
+      "wrap": true
+    },
+    "eval": {
+      "checkpoint_path": "",
+      "compute_generative_perplexity": true,
+      "compute_perplexity_on_sanity": false,
+      "disable_ema": false,
+      "ema_decay": null,
+      "gen_ppl_eval_model_name_or_path": "gpt2-large",
+      "generate_samples": true,
+      "generated_samples_path": "/home/david3684/projects/flm/hf_upload/samples.json",
+      "perplexity_batch_size": 8
+    },
+    "is_di4c": false,
+    "is_di4c_deterministic": true,
+    "loader": {
+      "batch_size": 32,
+      "eval_batch_size": 32,
+      "eval_global_batch_size": 512,
+      "global_batch_size": 512,
+      "num_workers": 4,
+      "pin_memory": true
+    },
+    "lr_scheduler": {
+      "_target_": "transformers.get_constant_schedule_with_warmup",
+      "num_warmup_steps": 2500
+    },
+    "mode": "train",
+    "model": {
+      "cond_dim": 128,
+      "dropout": 0.1,
+      "hidden_size": 768,
+      "length": 1024,
+      "n_blocks": 12,
+      "n_heads": 12,
+      "name": "small",
+      "scale_by_sigma": true,
+      "tie_word_embeddings": false,
+      "type": "ddit",
+      "vocab_lookup": true
+    },
+    "noise": {
+      "denoiser_latent_conditioning": -1,
+      "eps": 0,
+      "freeze_decoder": false,
+      "freeze_encoder": false,
+      "parameterization": "log-linear",
+      "type": "log-linear"
+    },
+    "optim": {
+      "beta1": 0.9,
+      "beta2": 0.999,
+      "eps": 1e-08,
+      "ln_tune": "none",
+      "lr": 0.0003,
+      "use_muon": false,
+      "weight_decay": 0
+    },
+    "prior": {
+      "latent_height": 0,
+      "latent_width": 0,
+      "type": "none"
+    },
+    "sampling": {
+      "argmax_correction": false,
+      "autoguidance_weight": 0.0,
+      "hard_start": false,
+      "noise_removal": "uniform_alpha",
+      "num_reflow_samples": 50000,
+      "num_sample_batches": 1,
+      "num_sample_log": 2,
+      "num_strides": 1,
+      "p_nucleus": 1.0,
+      "predictor": "ancestral",
+      "semi_ar": false,
+      "solver": "euler",
+      "steps": [
+        1024
+      ],
+      "stride_length": 1,
+      "tau_log10": -1.0,
+      "temperature": 1.0,
+      "use_float64": true
+    },
+    "seed": 1,
+    "strategy": {
+      "_target_": "lightning.pytorch.strategies.DDPStrategy",
+      "find_unused_parameters": false
+    },
+    "trainer": {
+      "_target_": "lightning.Trainer",
+      "accelerator": "cuda",
+      "accumulate_grad_batches": 2,
+      "devices": 8,
+      "gradient_clip_val": 1.0,
+      "limit_train_batches": 1.0,
+      "limit_val_batches": 1.0,
+      "log_every_n_steps": 100,
+      "max_steps": 1000000,
+      "num_nodes": 1,
+      "num_sanity_val_steps": 2,
+      "precision": "bf16",
+      "val_check_interval": 20000
+    },
+    "training": {
+      "antithetic_sampling": true,
+      "change_of_variables": false,
+      "double_ema": 0.0,
+      "ema": "0.9999, 0.999, 0.99, 0.95",
+      "finetune_path": "",
+      "importance_sampling": false,
+      "load_ema": false,
+      "loss_precision": "bf16",
+      "loss_type": "flow",
+      "loss_weighting": false,
+      "not_sampling_t": false,
+      "pred_type": "x0",
+      "sampling_eps": 0.001,
+      "t_curriculum_max": 1.0,
+      "t_curriculum_min": 0.0,
+      "t_curriculum_steps": 100000,
+      "torch_compile_mode": "default",
+      "use_torch_compile": true
+    },
+    "wandb": {
+      "group": null,
+      "id": "abcde2",
+      "job_type": null,
+      "name": "owt_ce_full_abcde2",
+      "notes": "DOS",
+      "project": "owt_full",
+      "tags": [
+        "openwebtext-train",
+        "dos"
+      ]
+    }
+  },
+  "vocab_size": 50258
+}

checkpoints/quarantine/20260821T024736Z/official/fmlm/owt/README.md

@@ -0,0 +1,47 @@
+---
+tags:
+- model_hub_mixin
+- pytorch_model_hub_mixin
+license: apache-2.0
+datasets:
+- Skylion007/openwebtext
+language:
+- en
+base_model:
+- david3684/FLM-B-OWT
+---
+  
+  ## FMLM (OWT)
+
+  For more information and implementation details, visit our github repository: [FLM](https://github.com/david3684/flm).
+
+  ## Model Details
+
+  The model has a context length of 1024, distilled from 1M steps trained FLM checkpoint trained on the OpenWebText (OWT) dataset.
+
+  For more details, please see our paper: 
+  - [arXiv:2602.16813](https://arxiv.org/abs/2602.16813)
+  - [Huggingface Papers](https://huggingface.co/papers/2602.16813)
+
+  ## Citation
+
+  Please cite our work using the bibtex below:
+
+  **BibTeX:**
+
+
+  ```bibtex
+  @article{lee2026flow,
+      title={Flow Map Language Models: One-step Language Modeling via Continuous Denoising},
+      author={Chanhyuk Lee and Jaehoon Yoo and Manan Agarwal
+              and Sheel Shah and Jerry Huang
+              and Aditi Raghunathan and Seunghoon Hong
+              and Nicholas M. Boffi and Jinwoo Kim},
+      journal={arXiv preprint arXiv:2602.16813},
+      year={2026},
+  }
+  ```
+
+  ## Model Card Contact
+
+  Chanhyuk Lee (chanhyuk-lee@kaist.ac.kr)

checkpoints/quarantine/20260821T024736Z/official/fmlm/owt/config.json

@@ -0,0 +1,223 @@
+{
+  "config": {
+    "algo": {
+      "T": 0,
+      "add_boundary": "fixed",
+      "backbone": "dit",
+      "backprop_entmax_temp_lsd": false,
+      "boundary_prob": 32,
+      "causal_attention": false,
+      "diagonal_fraction": 0.5,
+      "distillation_method": "PSD",
+      "double_temb": true,
+      "entmax_temp_lsd": 2.0,
+      "initialize_student_from_teacher": true,
+      "learnable_loss_weighting": false,
+      "loss_scale": "none",
+      "loss_type": null,
+      "name": "fmlm",
+      "offdiagonal_sampling": "uniform_diff",
+      "parameterization": "mean",
+      "pred_log_interval": 2000,
+      "rescale_offdiag_loss_psd": false,
+      "subs_masking": false,
+      "t_max": 1.0,
+      "t_min": 0.0,
+      "teacher_path": "/home/irteam/ours_owt_flow_ce_1m.ckpt",
+      "time_condition": "alpha_t",
+      "time_conditioning": true,
+      "use_discrete_schedule": true,
+      "use_ema_for_psd_target": false,
+      "use_mse_loss_psd": false,
+      "use_simulated_entropy": false,
+      "use_teacher_for_D_t_lsd": false
+    },
+    "callbacks": {
+      "checkpoint_every_n_steps": {
+        "_target_": "lightning.pytorch.callbacks.ModelCheckpoint",
+        "auto_insert_metric_name": false,
+        "dirpath": "/home/david3684/projects/flm/hf_upload/checkpoints",
+        "every_n_train_steps": 5000,
+        "save_last": false,
+        "save_top_k": -1,
+        "verbose": true
+      },
+      "checkpoint_monitor": {
+        "_target_": "lightning.pytorch.callbacks.ModelCheckpoint",
+        "auto_insert_metric_name": false,
+        "dirpath": "/home/david3684/projects/flm/hf_upload/checkpoints",
+        "filename": "best_nll",
+        "mode": "min",
+        "monitor": "val/nll",
+        "save_last": true,
+        "save_top_k": 1,
+        "verbose": true
+      },
+      "learning_rate_monitor": {
+        "_target_": "lightning.pytorch.callbacks.LearningRateMonitor",
+        "logging_interval": "step"
+      }
+    },
+    "checkpointing": {
+      "resume_ckpt_path": "/home/david3684/projects/flm/hf_upload/checkpoints/last.ckpt",
+      "resume_from_ckpt": false,
+      "save_dir": "/home/david3684/projects/flm/hf_upload"
+    },
+    "data": {
+      "cache_dir": "/home/irteam/workspace/datasets/openwebtext",
+      "insert_train_eos": true,
+      "insert_valid_eos": true,
+      "streaming": false,
+      "tokenizer_name_or_path": "gpt2",
+      "train": "openwebtext-train",
+      "valid": "openwebtext-valid",
+      "wrap": true
+    },
+    "eval": {
+      "checkpoint_path": "",
+      "compute_generative_perplexity": true,
+      "compute_perplexity_on_sanity": false,
+      "disable_ema": false,
+      "ema_decay": null,
+      "gen_ppl_eval_model_name_or_path": "gpt2-large",
+      "generate_samples": true,
+      "generated_samples_path": "/home/david3684/projects/flm/hf_upload/samples.json",
+      "perplexity_batch_size": 8
+    },
+    "is_di4c": false,
+    "is_di4c_deterministic": true,
+    "loader": {
+      "batch_size": 16,
+      "eval_batch_size": 16,
+      "eval_global_batch_size": 128,
+      "global_batch_size": 128,
+      "num_workers": 4,
+      "pin_memory": true
+    },
+    "lr_scheduler": {
+      "_target_": "transformers.get_constant_schedule_with_warmup",
+      "num_warmup_steps": 2500
+    },
+    "mode": "train",
+    "model": {
+      "cond_dim": 128,
+      "dropout": 0.1,
+      "hidden_size": 768,
+      "l2_normalize_embeddings": false,
+      "length": 1024,
+      "n_blocks": 12,
+      "n_heads": 12,
+      "name": "small",
+      "scale_by_sigma": true,
+      "tie_word_embeddings": false,
+      "type": "ddit",
+      "vocab_lookup": true
+    },
+    "noise": {
+      "denoiser_latent_conditioning": -1,
+      "eps": 0,
+      "freeze_decoder": false,
+      "freeze_encoder": false,
+      "parameterization": "log-linear",
+      "type": "log-linear"
+    },
+    "optim": {
+      "beta1": 0.9,
+      "beta2": 0.95,
+      "eps": 1e-08,
+      "ln_tune": "none",
+      "lr": 0.0003,
+      "muon_lr": 0.03,
+      "muon_momentum": 0.95,
+      "muon_weight_decay": 0.01,
+      "use_muon": false,
+      "weight_decay": 0
+    },
+    "prior": {
+      "latent_height": 0,
+      "latent_width": 0,
+      "type": "none"
+    },
+    "sampling": {
+      "argmax_correction": false,
+      "autoguidance_weight": 0.0,
+      "hard_start": false,
+      "noise_removal": "uniform_alpha",
+      "num_reflow_samples": 50000,
+      "num_sample_batches": 1,
+      "num_sample_log": 2,
+      "num_strides": 1,
+      "p_nucleus": 1.0,
+      "predictor": "ancestral",
+      "semi_ar": false,
+      "solver": "euler",
+      "steps": [
+        1,
+        2,
+        4,
+        8,
+        16,
+        32,
+        64,
+        128
+      ],
+      "stride_length": 1,
+      "tau_log10": -1.0,
+      "temperature": 1.0,
+      "use_float64": true
+    },
+    "seed": 1,
+    "strategy": {
+      "_target_": "lightning.pytorch.strategies.DDPStrategy",
+      "find_unused_parameters": false
+    },
+    "trainer": {
+      "_target_": "lightning.Trainer",
+      "accelerator": "cuda",
+      "accumulate_grad_batches": 1,
+      "devices": 8,
+      "gradient_clip_val": 1.0,
+      "limit_train_batches": 1.0,
+      "limit_val_batches": 10,
+      "log_every_n_steps": 100,
+      "max_steps": 1000000,
+      "num_nodes": 1,
+      "num_sanity_val_steps": 2,
+      "precision": "bf16",
+      "val_check_interval": 10000
+    },
+    "training": {
+      "antithetic_sampling": true,
+      "change_of_variables": false,
+      "double_ema": 0.0,
+      "ema": "0.9999, 0.999, 0.99, 0.95",
+      "finetune_path": "",
+      "importance_sampling": false,
+      "load_ema": false,
+      "loss_precision": "bf16",
+      "loss_type": "sd",
+      "loss_weighting": false,
+      "not_sampling_t": false,
+      "pred_type": "x0",
+      "sampling_eps": 0.001,
+      "t_curriculum_max": 1.0,
+      "t_curriculum_min": 0.0,
+      "t_curriculum_steps": 100000,
+      "torch_compile_mode": "default",
+      "use_torch_compile": true
+    },
+    "wandb": {
+      "group": null,
+      "id": "9efb5cfb",
+      "job_type": null,
+      "name": "owt_PSD_teacher_distill_LR_3e-4_DF_0.5_BP_32_SE_False_IT_True_EDMW_False_EMA_False_beta2_0.95_9efb5cfb",
+      "notes": "DOS",
+      "project": "owt_full",
+      "tags": [
+        "openwebtext-train",
+        "dos_denoiser"
+      ]
+    }
+  },
+  "vocab_size": 50258
+}

checkpoints/quarantine/20260821T024744Z/official/langflow/owt/config.json

@@ -0,0 +1,29 @@
+{
+  "_name_or_path": "chumengl/langflow-owt",
+  "architectures": [
+    "LangFlow"
+  ],
+  "auto_map": {
+    "AutoConfig": "config.LangFlowConfig",
+    "AutoModelForMaskedLM": "model.LangFlow"
+  },
+  "model_type": "LangFlow",
+  "vocab_size": 50257,
+  "hidden_size": 768,
+  "cond_dim": 128,
+  "n_blocks": 12,
+  "n_heads": 12,
+  "dropout": 0.1,
+  "model_length": 1024,
+  "use_normalized_embedding": true,
+  "embedding_norm_method": "layernorm",
+  "self_conditioning": true,
+  "use_bias": true,
+  "gumbel_loc": 4.723,
+  "gumbel_scale": 0.852,
+  "gumbel_cutoff": 1e-5,
+  "gumbel_entropy": 7.02,
+  "return_dict": true,
+  "torch_dtype": "float32",
+  "transformers_version": "4.38.2"
+}

checkpoints/quarantine/20260821T024744Z/official/langflow/owt/config.py

@@ -0,0 +1,63 @@
+"""HuggingFace configuration class for LangFlow."""
+
+import transformers
+
+
+class LangFlowConfig(transformers.PretrainedConfig):
+    """HuggingFace configuration class for LangFlow.
+    
+    LangFlow is a continuous diffusion language model that operates in embedding space.
+    It uses a DiT (Diffusion Transformer) backbone with adaptive layer normalization.
+    
+    Key features:
+        - Continuous diffusion in embedding space
+        - Self-conditioning: uses previous predictions as additional input
+        - Bias (preconditioning): skip connection for improved training
+        - Normalized embeddings: layernorm on embedding vectors
+        - Learnable Gumbel proposal for gamma (log-SNR) sampling
+    """
+    model_type = "LangFlow"
+
+    def __init__(
+        self,
+        vocab_size: int = 50257,
+        hidden_size: int = 768,
+        cond_dim: int = 128,
+        n_blocks: int = 12,
+        n_heads: int = 12,
+        dropout: float = 0.1,
+        model_length: int = 1024,
+        # Embedding normalization
+        use_normalized_embedding: bool = True,
+        embedding_norm_method: str = "layernorm",
+        # Self-conditioning
+        self_conditioning: bool = True,
+        # Bias (preconditioning) - always enabled for inference
+        use_bias: bool = True,
+        # Gumbel proposal parameters (learnable)
+        gumbel_loc: float = 4.723,
+        gumbel_scale: float = 0.852,
+        gumbel_cutoff: float = 1e-5,
+        gumbel_entropy: float = 7.02,
+        **kwargs
+    ):
+        super().__init__(**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.model_length = model_length
+        # Embedding normalization
+        self.use_normalized_embedding = use_normalized_embedding
+        self.embedding_norm_method = embedding_norm_method
+        # Self-conditioning
+        self.self_conditioning = self_conditioning
+        # Bias (preconditioning)
+        self.use_bias = use_bias
+        # Gumbel proposal parameters
+        self.gumbel_loc = gumbel_loc
+        self.gumbel_scale = gumbel_scale
+        self.gumbel_cutoff = gumbel_cutoff
+        self.gumbel_entropy = gumbel_entropy

checkpoints/quarantine/20260825T100043Z/official/flm/lm1b/README.md

@@ -0,0 +1,43 @@
+---
+tags:
+- model_hub_mixin
+- pytorch_model_hub_mixin
+language:
+- en
+license: apache-2.0
+---
+  
+  ## FLM (LM1B)
+
+  For more information and implementation details, visit our github repository: [FLM](https://github.com/david3684/flm).
+
+  ## Model Details
+
+  The model has a context length of 128, trained for 1M steps on the One Billion Words (LM1B) dataset.
+
+  For more details, please see our paper:
+  - [arXiv:2602.16813](https://arxiv.org/abs/2602.16813)
+  - [Huggingface Papers](https://huggingface.co/papers/2602.16813)
+    
+  ## Citation
+
+  Please cite our work using the bibtex below:
+
+  **BibTeX:**
+
+
+  ```bibtex
+  @article{lee2026flow,
+      title={Flow Map Language Models: One-step Language Modeling via Continuous Denoising},
+      author={Chanhyuk Lee and Jaehoon Yoo and Manan Agarwal
+              and Sheel Shah and Jerry Huang
+              and Aditi Raghunathan and Seunghoon Hong
+              and Nicholas M. Boffi and Jinwoo Kim},
+      journal={arXiv preprint arXiv:2602.16813},
+      year={2026},
+  }
+  ```
+
+  ## Model Card Contact
+
+  Chanhyuk Lee (chanhyuk-lee@kaist.ac.kr)

checkpoints/quarantine/20260825T100259Z/official/duo_dcd/owt/config.py

@@ -0,0 +1,29 @@
+import transformers
+
+
+class DUOConfig(transformers.PretrainedConfig):
+  """Hugging Face configuration class for DUO."""
+  model_type = 'DUO'
+
+  def __init__(
+    self,
+    vocab_size: int = 50258,
+    model_length: int = 1024,
+    causal: bool = False,
+    hidden_dim: int = 768,
+    cond_dim: int = 129,
+    n_blocks: int = 12,
+    n_heads: int = 12,
+    dropout: float = 0.1,
+    var_min: bool = True,
+    ** kwargs):
+    super().__init__(**kwargs)
+    self.causal = causal
+    self.vocab_size = vocab_size
+    self.model_length = model_length
+    self.hidden_dim = hidden_dim
+    self.cond_dim = cond_dim
+    self.n_blocks = n_blocks
+    self.n_heads = n_heads
+    self.dropout = dropout
+    self.var_min = var_min

checkpoints/quarantine/20260825T100259Z/official/duo_dcd/owt/model.py

@@ -0,0 +1,601 @@
+import math
+import typing
+
+import einops
+import flash_attn
+import flash_attn.layers.rotary
+import huggingface_hub
+import omegaconf
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import transformers
+
+from .config import DUOConfig
+
+# Flags required to enable jit fusion kernels
+torch._C._jit_set_profiling_mode(False)
+torch._C._jit_set_profiling_executor(False)
+torch._C._jit_override_can_fuse_on_cpu(True)
+torch._C._jit_override_can_fuse_on_gpu(True)
+
+
+def bias_dropout_add_scale(
+    x: torch.Tensor,
+    bias: typing.Optional[torch.Tensor],
+    scale: torch.Tensor,
+    residual: typing.Optional[torch.Tensor],
+    prob: float,
+    training: bool) -> torch.Tensor:
+  if bias is not None:
+    out = scale * F.dropout(x + bias, p=prob, training=training)
+  else:
+    out = scale * F.dropout(x, p=prob, training=training)
+
+  if residual is not None:
+    out = residual + out
+  return out
+
+
+def get_bias_dropout_add_scale(training):
+  def _bias_dropout_add(x, bias, scale, residual, prob):
+    return bias_dropout_add_scale(
+      x, bias, scale, residual, prob, training)
+
+  return _bias_dropout_add
+
+
+# function overload
+def modulate(x: torch.Tensor,
+             shift: torch.Tensor,
+             scale: torch.Tensor) -> torch.Tensor:
+  return x * (1 + scale) + shift
+
+
+@torch.jit.script
+def bias_dropout_add_scale_fused_train(
+    x: torch.Tensor,
+    bias: typing.Optional[torch.Tensor],
+    scale: torch.Tensor,
+    residual: typing.Optional[torch.Tensor],
+    prob: float) -> torch.Tensor:
+  return bias_dropout_add_scale(
+    x, bias, scale, residual, prob, True)
+
+
+@torch.jit.script
+def bias_dropout_add_scale_fused_inference(
+    x: torch.Tensor,
+    bias: typing.Optional[torch.Tensor],
+    scale: torch.Tensor,
+    residual: typing.Optional[torch.Tensor],
+    prob: float) -> torch.Tensor:
+  return bias_dropout_add_scale(
+    x, bias, scale, residual, prob, False)
+
+
+@torch.jit.script
+def modulate_fused(x: torch.Tensor,
+                   shift: torch.Tensor,
+                   scale: torch.Tensor) -> torch.Tensor:
+  return modulate(x, shift, scale)
+
+
+class Rotary(torch.nn.Module):
+  def __init__(self, dim, base=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, seq_dim=1):
+    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.)
+      self.sin_cached[:,:,2,:,:].fill_(0.)
+
+    return self.cos_cached, self.sin_cached
+
+
+def rotate_half(x):
+  x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :]
+  return torch.cat((-x2, x1), dim=-1)
+
+
+def split_and_apply_rotary_pos_emb(qkv, rotary_cos_sin):
+  with torch.cuda.amp.autocast(enabled=False):
+    cos, sin = rotary_cos_sin
+    cos = cos.to(qkv.dtype)
+    sin = sin.to(qkv.dtype)
+    cos = cos[0,:,0,0,:cos.shape[-1]//2]
+    sin = sin[0,:,0,0,:sin.shape[-1]//2]
+    q, k, v = qkv.chunk(3, dim=2)
+    q = flash_attn.layers.rotary.apply_rotary_emb_torch(
+      q.squeeze(dim=2), cos, sin)
+    k = flash_attn.layers.rotary.apply_rotary_emb_torch(
+      k.squeeze(dim=2), cos, sin)
+    v = v.squeeze(dim=2)
+  return q, k, v
+
+
+def apply_rotary_pos_emb(qkv, cos, sin):
+  cos = cos[0,:,0,0,:cos.shape[-1]//2]
+  sin = sin[0,:,0,0,:sin.shape[-1]//2]
+  return flash_attn.layers.rotary.apply_rotary_emb_qkv_(qkv, cos, sin)
+
+
+def regular_attention_multi_headed(q, k, v):
+  # Assuming qkv is a tensor with shape [batch, seq_len, 3, num_heads, head_dim]
+  # where the 3 represents Q, K, V packed in that order
+  attention_output = F.scaled_dot_product_attention(
+    query=q.transpose(1, 2),
+    key=k.transpose(1, 2),
+    value=v.transpose(1, 2),
+    attn_mask=None,
+    dropout_p=0.0,
+    is_causal=False)
+  # [batch_size, seq_len, num_heads, head_dim]
+  attention_output = attention_output.transpose(1, 2)
+  return einops.rearrange(attention_output, 'b s h d -> b s (h d)')
+
+
+#################################################################################
+#                                  Layers                                       #
+#################################################################################
+class LayerNorm(nn.Module):
+  def __init__(self, dim):
+    super().__init__()
+    self.weight = nn.Parameter(torch.ones([dim]))
+    self.dim = dim
+  def forward(self, x):
+    with torch.cuda.amp.autocast(enabled=False):
+      x = F.layer_norm(x.float(), [self.dim])
+    return x * self.weight[None, None, :]
+
+
+def residual_linear(x, W, x_skip, residual_scale):
+  """x_skip + residual_scale * W @ x"""
+  dim_out, dim_in = W.shape[0], W.shape[1]
+  return torch.addmm(
+    x_skip.view(-1, dim_out),
+    x.view(-1, dim_in),
+    W.T,
+    alpha=residual_scale).view(*x.shape[:-1], dim_out)
+
+
+#################################################################################
+#               Embedding Layers for Timesteps and Class Labels                 #
+#################################################################################
+class TimestepEmbedder(nn.Module):
+  """
+  Embeds scalar timesteps into vector representations.
+  """
+  def __init__(self, hidden_size, frequency_embedding_size=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(t, dim, max_period=10000):
+    """
+    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.
+    """
+    # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
+    half = dim // 2
+    freqs = torch.exp(
+      - math.log(max_period)
+      * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device)
+      / half)
+    args = t[:, 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, t):
+    t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+    t_emb = self.mlp(t_freq)
+    return t_emb
+
+
+class LabelEmbedder(nn.Module):
+  """Embeds class labels into vector representations.
+  
+  Also handles label dropout for classifier-free guidance.
+  """
+  def __init__(self, num_classes, cond_size):
+    super().__init__()
+    self.embedding_table = nn.Embedding(num_classes + 1, cond_size)
+    self.num_classes = num_classes
+
+    # TODO think of initializing with 0.02 std deviation like in original DiT paper
+
+  def forward(self, labels):
+    embeddings = self.embedding_table(labels)
+    return embeddings
+    
+
+#################################################################################
+#                                 Core Model                                    #
+#################################################################################
+
+class DDiTBlockCausal(nn.Module):
+  def __init__(self, dim, n_heads, mlp_ratio=4, dropout=0.1):
+    super().__init__()
+    self.n_heads = n_heads
+
+    self.norm1 = LayerNorm(dim)
+    self.attn_qkv = nn.Linear(dim, 3 * dim, bias=False)
+    self.attn_out = nn.Linear(dim, dim, bias=False)
+    self.dropout1 = nn.Dropout(dropout)
+
+    self.norm2 = LayerNorm(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.dropout2 = nn.Dropout(dropout)
+    self.dropout = dropout
+
+  def _get_bias_dropout_scale(self):
+    if self.training:
+      return bias_dropout_add_scale_fused_train
+    else:
+      return bias_dropout_add_scale_fused_inference
+
+  def forward(self, x, rotary_cos_sin, **kwargs):
+    del kwargs
+    batch_size, seq_len = x.shape[0], x.shape[1]
+
+    bias_dropout_scale_fn = self._get_bias_dropout_scale()
+
+    # attention operation
+    x_skip = x
+    x = self.norm1(x)
+
+    qkv = self.attn_qkv(x)
+    qkv = einops.rearrange(
+      qkv,
+      'b s (three h d) -> b s three h d',
+      three=3,
+      h=self.n_heads)
+    with torch.cuda.amp.autocast(enabled=False):
+      cos, sin = rotary_cos_sin
+      qkv = apply_rotary_pos_emb(
+        qkv, cos.to(qkv.dtype), sin.to(qkv.dtype)
+      )
+    qkv = einops.rearrange(qkv, 'b s ... -> (b s) ...')
+    cu_seqlens = torch.arange(
+      0, (batch_size + 1) * seq_len,
+      step=seq_len, dtype=torch.int32, device=qkv.device)
+    x = flash_attn.flash_attn_interface.flash_attn_varlen_qkvpacked_func(
+      qkv, cu_seqlens, seq_len, 0.0, causal=True)
+
+    x = einops.rearrange(x, '(b s) h d -> b s (h d)',
+                         b=batch_size)
+
+    scale = torch.ones(1, device=x.device, dtype=x.dtype)
+    x = bias_dropout_scale_fn(
+      self.attn_out(x), None, scale, x_skip, self.dropout)
+
+    # mlp operation
+    x = bias_dropout_scale_fn(
+      self.mlp(self.norm2(x)), None, scale, x, self.dropout)
+    return x
+
+
+
+class DDiTBlock(nn.Module):
+  def __init__(self, dim, n_heads, adaLN,
+               cond_dim=None, mlp_ratio=4,
+               dropout=0.1):
+    super().__init__()
+    self.n_heads = n_heads
+    self.adaLN = adaLN
+
+    self.norm1 = LayerNorm(dim)
+    self.attn_qkv = nn.Linear(dim, 3 * dim, bias=False)
+    self.attn_out = nn.Linear(dim, dim, bias=False)
+    self.dropout1 = nn.Dropout(dropout)
+
+    self.norm2 = LayerNorm(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.dropout2 = nn.Dropout(dropout)
+    self.dropout = dropout
+
+    if self.adaLN:
+      self.adaLN_modulation = nn.Linear(cond_dim, 6 * dim)
+      self.adaLN_modulation.weight.data.zero_()
+      self.adaLN_modulation.bias.data.zero_()
+
+
+  def _get_bias_dropout_scale(self):
+    if self.training:
+      return bias_dropout_add_scale_fused_train
+    else:
+      return bias_dropout_add_scale_fused_inference
+
+
+  def forward(self, x, rotary_cos_sin, c=None):
+
+    bias_dropout_scale_fn = self._get_bias_dropout_scale()
+
+    x_skip = x
+    x = self.norm1(x)
+
+    if self.adaLN:
+      # self.adaLN_modulation(c): (128, 1536)
+      # self.adaLN_modulation(c)[:, None]: (128, 1, 1536)
+      # "" .chunk(6, dim=2) returns 6 tuples of shapes (128, 1, 256)
+      (shift_msa, scale_msa, gate_msa, shift_mlp,
+       scale_mlp, gate_mlp) = self.adaLN_modulation(c)[:, None].chunk(6, dim=2)
+      x = modulate_fused(x, shift_msa, scale_msa)
+
+    qkv = einops.rearrange(
+      self.attn_qkv(x),
+      'b s (three h d) -> b s three h d',
+      three=3,
+      h=self.n_heads)
+    q, k, v = split_and_apply_rotary_pos_emb(qkv, rotary_cos_sin)
+    
+    x = regular_attention_multi_headed(q, k, v)
+
+    if self.adaLN:
+      x = bias_dropout_scale_fn(self.attn_out(x),
+                                None,
+                                gate_msa,
+                                x_skip,
+                                self.dropout)
+      x = bias_dropout_scale_fn(
+        self.mlp(modulate_fused(
+          self.norm2(x), shift_mlp, scale_mlp)),
+        None, gate_mlp, x, self.dropout)
+    else:
+      scale = torch.ones(1, device=x.device, dtype=x.dtype)
+      x = bias_dropout_scale_fn(
+        self.attn_out(x), None, scale, x_skip, self.dropout)
+      x = bias_dropout_scale_fn(
+        self.mlp(self.norm2(x)), None, scale, x, self.dropout)
+    return x
+
+
+class EmbeddingLayer(nn.Module):
+  def __init__(self, dim, vocab_dim):
+    super().__init__()
+    self.embedding = nn.Parameter(torch.empty((vocab_dim, dim)))
+    torch.nn.init.kaiming_uniform_(self.embedding, a=math.sqrt(5))
+
+  def forward(self, x):
+    if x.ndim == 2:
+      return self.embedding[x]
+    assert x.ndim == 3
+    return torch.einsum(
+      "blv,ve->ble",
+      torch.nn.functional.softmax(x, dim=-1).float(),
+      self.embedding.float()).to(x.dtype)
+
+
+class DDiTFinalLayer(nn.Module):
+  def __init__(self, hidden_size, out_channels, cond_dim,
+               adaLN):
+    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 = adaLN
+    if self.adaLN:
+      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, c):
+    x = self.norm_final(x)
+    if self.adaLN:
+      shift, scale = self.adaLN_modulation(c)[:, None].chunk(2, dim=2)
+      x = modulate_fused(x, shift, scale)
+    x = self.linear(x)
+    return x
+
+
+class DIT(nn.Module, huggingface_hub.PyTorchModelHubMixin):
+  def __init__(self, config, vocab_size: int):
+    super().__init__()
+    if type(config) == dict:
+      config = omegaconf.OmegaConf.create(config)
+    self.causal = config.algo.causal_attention
+    self.adaLN = not self.causal
+    self.config = config
+    self.vocab_size = vocab_size
+    dim = config.model.hidden_size
+    cond_dim = config.model.cond_dim
+    self.vocab_embed = EmbeddingLayer(dim, vocab_size)
+    if not self.causal:
+      self.sigma_map = TimestepEmbedder(cond_dim)
+    self.rotary_emb = Rotary(dim // config.model.n_heads)
+
+    blocks = []
+    for _ in range(config.model.n_blocks):
+      if self.causal:
+        block = DDiTBlockCausal(
+          dim=dim,
+          n_heads=config.model.n_heads,
+          dropout=config.model.dropout)
+      else:
+        block = DDiTBlock(
+          dim=dim,
+          n_heads=config.model.n_heads,
+          cond_dim=cond_dim,
+          adaLN=self.adaLN,
+          dropout=config.model.dropout)
+      blocks.append(block)
+    self.blocks = nn.ModuleList(blocks)
+
+    self.output_layer = DDiTFinalLayer(
+      hidden_size=dim,
+      out_channels=vocab_size,
+      cond_dim=cond_dim,
+      adaLN=self.adaLN)
+    self.scale_by_sigma = config.model.scale_by_sigma
+
+  def _get_bias_dropout_scale(self):
+    if self.training:
+      return bias_dropout_add_scale_fused_train
+    else:
+      return  bias_dropout_add_scale_fused_inference
+
+  def forward(self, x, sigma):
+    x = self.vocab_embed(x)
+    if self.causal:
+      t_cond = None
+    else:
+      t_cond = F.silu(self.sigma_map(sigma))
+
+    rotary_cos_sin = self.rotary_emb(x).to(x.device)
+
+    with torch.cuda.amp.autocast(dtype=torch.bfloat16):
+      for i in range(len(self.blocks)):
+        x = self.blocks[i](x, rotary_cos_sin, c=t_cond)
+      x = self.output_layer(x, c=t_cond)
+
+    return x
+
+
+
+class HFDIT(torch.nn.Module):
+  def __init__(self, config):
+    super().__init__()
+    self.causal = config.causal
+    self.adaLN = not self.causal
+    self.vocab_size = config.vocab_size
+    dim = config.hidden_dim
+    cond_dim = config.cond_dim
+    self.vocab_embed = EmbeddingLayer(dim, self.vocab_size)
+    if not self.causal:
+      self.sigma_map = TimestepEmbedder(cond_dim)
+    self.rotary_emb = Rotary(dim // config.n_heads)
+
+    blocks = []
+    for _ in range(config.n_blocks):
+      if self.causal:
+        block = DDiTBlockCausal(
+          dim=dim,
+          n_heads=config.n_heads,
+          dropout=config.dropout)
+      else:
+        block = DDiTBlock(
+          dim=dim,
+          n_heads=config.n_heads,
+          cond_dim=cond_dim,
+          adaLN=self.adaLN,
+          dropout=config.dropout)
+      blocks.append(block)
+    self.blocks = torch.nn.ModuleList(blocks)
+
+    self.output_layer = DDiTFinalLayer(
+      hidden_size=dim,
+      out_channels=self.vocab_size,
+      cond_dim=cond_dim,
+      adaLN=self.adaLN)
+
+  def _get_bias_dropout_scale(self):
+    if self.training:
+      return bias_dropout_add_scale_fused_train
+    else:
+      return  bias_dropout_add_scale_fused_inference
+
+  def forward(self, x, sigma, output_hidden_states=False):
+    all_hidden_states = []
+    x = self.vocab_embed(x)
+    if output_hidden_states:
+      all_hidden_states.append(x)
+    if self.causal:
+      t_cond = None
+    else:
+      t_cond = F.silu(self.sigma_map(sigma))
+
+    rotary_cos_sin = self.rotary_emb(x)
+    with torch.cuda.amp.autocast(
+      dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float32):
+      for i in range(len(self.blocks)):
+        x = self.blocks[i](x, rotary_cos_sin, c=t_cond)
+        if output_hidden_states:
+          all_hidden_states.append(x)
+      x = self.output_layer(x, c=t_cond)
+    return x, all_hidden_states
+
+
+class DUO(transformers.PreTrainedModel):
+  """HF-compatible model."""
+  config_class = DUOConfig
+  base_model_prefix = 'duo'
+
+  def __init__(self, config: DUOConfig):
+    super().__init__(config)
+    self.config = config
+    self.backbone = HFDIT(config)
+
+  def reset_kv_cache(self):
+    for block in self.backbone.blocks:
+      block.kv_cache = None
+
+  def forward(
+      self,
+      input_ids: torch.LongTensor = None,
+      timesteps: torch.FloatTensor = None,
+      output_hidden_states: typing.Optional[bool] = None,
+      return_dict: typing.Optional[bool] = None,
+  ) -> typing.Union[
+    torch.Tensor, typing.Tuple,
+    transformers.modeling_outputs.MaskedLMOutput]:
+    """HF-compatible forward method."""
+    output_hidden_states = (
+      output_hidden_states
+      if output_hidden_states is not None
+      else self.config.output_hidden_states
+    )
+    return_dict = return_dict \
+      if return_dict is not None \
+      else self.config.use_return_dict
+    
+    logits, all_hidden_states = self.backbone(
+      x=input_ids,
+      sigma=timesteps,
+      output_hidden_states=output_hidden_states,
+    )
+    if return_dict:
+      return transformers.modeling_outputs.MaskedLMOutput(
+        logits=logits,
+        hidden_states=all_hidden_states if output_hidden_states else None,
+        loss=None
+      )
+    elif output_hidden_states:
+      return logits, all_hidden_states
+    else:
+      return logits

checkpoints/reference_reproduction/_invalid/old_di4c_shape_skip_20260825_095622/duo_di4c_lm1b/.hydra/config.yaml

@@ -0,0 +1,179 @@
+mode: train
+time_conditioning: false
+T: 1024
+compile: false
+eval_every: 5000000
+is_di4c: true
+is_teacher_di4c: false
+latent_bsize: 16
+alpha_t: sigmoid
+alpha_const: 0.1
+distil_delta: 0.02
+log_unif: false
+round: 7
+data_preprocess:
+  data_cache: /home/ruizzhao/diffusion_baseline/data/training_cache/di4c/lm1b
+  min_seq_len: -1
+  seq_len: 128
+  group_text: true
+  remove_text: true
+  num_seqs: -1
+  add_bos: false
+  add_eos: true
+  legacy_start_end_bos: true
+loader:
+  global_batch_size: 2
+  eval_global_batch_size: 2
+  batch_size: 1
+  eval_batch_size: 1
+  num_workers: ${eval:"len(__import__('os').sched_getaffinity(0))"}
+  pin_memory: true
+  persistent_workers: true
+training:
+  ema: 0.9999
+  antithetic_sampling: true
+  importance_sampling: false
+  sampling_eps: 1.0e-05
+  change_of_variables: false
+tokenizer:
+  name: /home/ruizzhao/diffusion_baseline/data/raw/huggingface/hub/models--bert-base-uncased/snapshots/86b5e0934494bd15c9632b12f734a8a67f723594
+optim:
+  name: adamw
+  weight_decay: 0.0
+  lr: 3.0e-05
+  beta1: 0.9
+  beta2: 0.999
+  eps: 1.0e-08
+trainer:
+  _target_: lightning.Trainer
+  accelerator: cuda
+  num_nodes: 1
+  devices: 1
+  accumulate_grad_batches: 2
+  gradient_clip_val: 1.0
+  precision: bf16-mixed
+  num_sanity_val_steps: 0
+  max_steps: 20000
+  log_every_n_steps: 10
+  limit_train_batches: 1.0
+  limit_val_batches: 1.0
+  val_check_interval: ${eval:${trainer.accumulate_grad_batches} * ${eval_every}}
+  check_val_every_n_epoch: null
+  benchmark: true
+wandb: null
+checkpointing:
+  save_dir: /home/ruizzhao/diffusion_baseline/checkpoints/reference_reproduction/duo_di4c/lm1b
+  resume_from_ckpt: true
+  resume_ckpt_path: /home/ruizzhao/diffusion_baseline/checkpoints/reference_reproduction/duo_di4c/lm1b/checkpoints/last.ckpt
+eval:
+  valid:
+    n_samples: 16
+    num_steps: 128
+  ppl_with_ar:
+    run: true
+    model: gpt2-large
+    batch_size: 8
+  cond_ppl:
+    run: true
+    model: gpt2-large
+    batch_size: 8
+  mauve:
+    run: false
+    num_rounds: 5
+    max_num_tokens: 100
+    batch_size: 8
+    scaling_factor: 5
+  lambada_openai:
+    run: false
+    batch_size: 32
+    num_samples: 32
+    from_ema: false
+    add_bos: true
+    add_eos: true
+callbacks:
+  checkpoint_every_n_steps:
+    _target_: lightning.pytorch.callbacks.ModelCheckpoint
+    save_top_k: -1
+    save_last: true
+    dirpath: ${checkpointing.save_dir}/checkpoints
+    verbose: true
+    auto_insert_metric_name: false
+    every_n_train_steps: 10000
+  learning_rate_monitor:
+    _target_: lightning.pytorch.callbacks.LearningRateMonitor
+    logging_interval: step
+  text_throughput:
+    _target_: callbacks.text_throughput.TextThroughputMonitor
+  grad_norm:
+    _target_: callbacks.gradient_norm.GradNormCallback
+data:
+  train: lm1b
+  valid: lm1b
+  test: openwebtext-test
+  tokenizer_name_or_path: gpt2
+  wrap: true
+  streaming: false
+model:
+  name: small
+  type: ddit-orig
+  hidden_size: 768
+  cond_dim: 128
+  length: 128
+  n_blocks: 12
+  n_heads: 12
+  scale_by_sigma: true
+  dropout: 0.1
+  tie_word_embeddings: false
+strategy:
+  _target_: lightning.pytorch.strategies.DDPStrategy
+  find_unused_parameters: true
+noise:
+  type: loglinear
+  sigma_min: 0.0001
+  sigma_max: 20
+lr_scheduler:
+  _target_: transformers.get_constant_schedule_with_warmup
+  num_warmup_steps: 2500
+parameterization:
+  name: multi-round-sdtt
+  log_loss_buckets: -1
+  checkpoint_path: /home/ruizzhao/diffusion_baseline/checkpoints/reference_reproduction/duo_di4c/lm1b/inputs/teacher_adapted.ckpt
+  start_from_hf: false
+  distill_mode: kl-bwd
+  num_distill_steps: 2
+  min_num_sampling_steps: 8
+  grow_dt_every: 10000
+  orig_num_sampling_steps: 1024
+  sampling_mode: ancestral
+  loss_precision: null
+  reset_optimizer_on_growth: false
+  use_ema_on_growth: false
+  sampling:
+    uncond:
+      run: true
+      num_samples: 1024
+      batch_size: 16
+      from_ema: false
+      num_steps: 16
+      seq_len: 1024
+      sampler: ancestral
+      cache_preds: false
+      add_bos: false
+      add_eos: false
+    cond_prefix:
+      run: true
+      num_samples: 256
+      batch_size: 16
+      from_ema: false
+      dataset: webtext
+      seq_len: 100
+      prefix_len: 50
+      num_cont_per_prefix: 5
+      min_seq_len: 1024
+      num_steps: 128
+      sampler: ancestral
+      cache_preds: false
+      add_bos: false
+      add_eos: false
+      calc_sbleu: true
+seed: 42

checkpoints/reference_reproduction/_invalid/old_di4c_shape_skip_20260825_095622/duo_di4c_lm1b/.hydra/hydra.yaml

@@ -0,0 +1,202 @@
+hydra:
+  run:
+    dir: /home/ruizzhao/diffusion_baseline/checkpoints/reference_reproduction/duo_di4c/lm1b
+  sweep:
+    dir: multirun/${now:%Y-%m-%d}/${now:%H-%M-%S}
+    subdir: ${hydra.job.num}
+  launcher:
+    _target_: hydra._internal.core_plugins.basic_launcher.BasicLauncher
+  sweeper:
+    _target_: hydra._internal.core_plugins.basic_sweeper.BasicSweeper
+    max_batch_size: null
+    params: null
+  help:
+    app_name: ${hydra.job.name}
+    header: '${hydra.help.app_name} is powered by Hydra.
+
+      '
+    footer: 'Powered by Hydra (https://hydra.cc)
+
+      Use --hydra-help to view Hydra specific help
+
+      '
+    template: '${hydra.help.header}
+
+      == Configuration groups ==
+
+      Compose your configuration from those groups (group=option)
+
+
+      $APP_CONFIG_GROUPS
+
+
+      == Config ==
+
+      Override anything in the config (foo.bar=value)
+
+
+      $CONFIG
+
+
+      ${hydra.help.footer}
+
+      '
+  hydra_help:
+    template: 'Hydra (${hydra.runtime.version})
+
+      See https://hydra.cc for more info.
+
+
+      == Flags ==
+
+      $FLAGS_HELP
+
+
+      == Configuration groups ==
+
+      Compose your configuration from those groups (For example, append hydra/job_logging=disabled
+      to command line)
+
+
+      $HYDRA_CONFIG_GROUPS
+
+
+      Use ''--cfg hydra'' to Show the Hydra config.
+
+      '
+    hydra_help: ???
+  hydra_logging:
+    version: 1
+    formatters:
+      simple:
+        format: '[%(asctime)s][HYDRA] %(message)s'
+    handlers:
+      console:
+        class: logging.StreamHandler
+        formatter: simple
+        stream: ext://sys.stdout
+    root:
+      level: INFO
+      handlers:
+      - console
+    loggers:
+      logging_example:
+        level: DEBUG
+    disable_existing_loggers: false
+  job_logging:
+    version: 1
+    formatters:
+      simple:
+        format: '[%(asctime)s][%(name)s][%(levelname)s] - %(message)s'
+    handlers:
+      console:
+        class: logging.StreamHandler
+        formatter: simple
+        stream: ext://sys.stdout
+      file:
+        class: logging.FileHandler
+        formatter: simple
+        filename: ${hydra.runtime.output_dir}/${hydra.job.name}.log
+    root:
+      level: INFO
+      handlers:
+      - console
+      - file
+    disable_existing_loggers: false
+  env: {}
+  mode: RUN
+  searchpath: []
+  callbacks: {}
+  output_subdir: .hydra
+  overrides:
+    hydra:
+    - hydra.run.dir=/home/ruizzhao/diffusion_baseline/checkpoints/reference_reproduction/duo_di4c/lm1b
+    - hydra.job.chdir=true
+    - hydra.mode=RUN
+    task:
+    - mode=train
+    - data.train=lm1b
+    - data.valid=lm1b
+    - tokenizer.name=/home/ruizzhao/diffusion_baseline/data/raw/huggingface/hub/models--bert-base-uncased/snapshots/86b5e0934494bd15c9632b12f734a8a67f723594
+    - data_preprocess.data_cache=/home/ruizzhao/diffusion_baseline/data/training_cache/di4c/lm1b
+    - data_preprocess.seq_len=128
+    - data_preprocess.group_text=true
+    - data_preprocess.remove_text=true
+    - data_preprocess.add_bos=false
+    - data_preprocess.add_eos=true
+    - data_preprocess.legacy_start_end_bos=true
+    - model=dit-orig-small
+    - model.length=128
+    - parameterization=multi-round-sdtt
+    - parameterization.checkpoint_path=/home/ruizzhao/diffusion_baseline/checkpoints/reference_reproduction/duo_di4c/lm1b/inputs/teacher_adapted.ckpt
+    - parameterization.start_from_hf=false
+    - parameterization.grow_dt_every=10000
+    - parameterization.orig_num_sampling_steps=1024
+    - loader.global_batch_size=2
+    - loader.eval_global_batch_size=2
+    - loader.batch_size=1
+    - loader.eval_batch_size=1
+    - trainer.devices=1
+    - trainer.num_nodes=1
+    - trainer.accumulate_grad_batches=2
+    - trainer.max_steps=20000
+    - trainer.precision=bf16-mixed
+    - optim.lr=3e-05
+    - lr_scheduler.num_warmup_steps=2500
+    - callbacks.checkpoint_every_n_steps.every_n_train_steps=10000
+    - checkpointing.save_dir=/home/ruizzhao/diffusion_baseline/checkpoints/reference_reproduction/duo_di4c/lm1b
+    - checkpointing.resume_from_ckpt=true
+    - checkpointing.resume_ckpt_path=/home/ruizzhao/diffusion_baseline/checkpoints/reference_reproduction/duo_di4c/lm1b/checkpoints/last.ckpt
+    - wandb=null
+    - +seed=42
+    - is_di4c=true
+    - is_teacher_di4c=false
+    - T=1024
+    - latent_bsize=16
+    - round=7
+  job:
+    name: main
+    chdir: true
+    override_dirname: +seed=42,T=1024,callbacks.checkpoint_every_n_steps.every_n_train_steps=10000,checkpointing.resume_ckpt_path=/home/ruizzhao/diffusion_baseline/checkpoints/reference_reproduction/duo_di4c/lm1b/checkpoints/last.ckpt,checkpointing.resume_from_ckpt=true,checkpointing.save_dir=/home/ruizzhao/diffusion_baseline/checkpoints/reference_reproduction/duo_di4c/lm1b,data.train=lm1b,data.valid=lm1b,data_preprocess.add_bos=false,data_preprocess.add_eos=true,data_preprocess.data_cache=/home/ruizzhao/diffusion_baseline/data/training_cache/di4c/lm1b,data_preprocess.group_text=true,data_preprocess.legacy_start_end_bos=true,data_preprocess.remove_text=true,data_preprocess.seq_len=128,is_di4c=true,is_teacher_di4c=false,latent_bsize=16,loader.batch_size=1,loader.eval_batch_size=1,loader.eval_global_batch_size=2,loader.global_batch_size=2,lr_scheduler.num_warmup_steps=2500,mode=train,model.length=128,model=dit-orig-small,optim.lr=3e-05,parameterization.checkpoint_path=/home/ruizzhao/diffusion_baseline/checkpoints/reference_reproduction/duo_di4c/lm1b/inputs/teacher_adapted.ckpt,parameterization.grow_dt_every=10000,parameterization.orig_num_sampling_steps=1024,parameterization.start_from_hf=false,parameterization=multi-round-sdtt,round=7,tokenizer.name=/home/ruizzhao/diffusion_baseline/data/raw/huggingface/hub/models--bert-base-uncased/snapshots/86b5e0934494bd15c9632b12f734a8a67f723594,trainer.accumulate_grad_batches=2,trainer.devices=1,trainer.max_steps=20000,trainer.num_nodes=1,trainer.precision=bf16-mixed,wandb=null
+    id: ???
+    num: ???
+    config_name: config
+    env_set: {}
+    env_copy: []
+    config:
+      override_dirname:
+        kv_sep: '='
+        item_sep: ','
+        exclude_keys: []
+  runtime:
+    version: 1.3.2
+    version_base: '1.3'
+    cwd: /home/ruizzhao/diffusion_baseline/upstreams/di4c
+    config_sources:
+    - path: hydra.conf
+      schema: pkg
+      provider: hydra
+    - path: /home/ruizzhao/diffusion_baseline/upstreams/di4c/sdtt/src/sdtt/configs
+      schema: file
+      provider: main
+    - path: ''
+      schema: structured
+      provider: schema
+    output_dir: /home/ruizzhao/diffusion_baseline/checkpoints/reference_reproduction/duo_di4c/lm1b
+    choices:
+      parameterization: multi-round-sdtt
+      lr_scheduler: constant_warmup
+      noise: loglinear
+      strategy: ddp
+      model: dit-orig-small
+      data: openwebtext-split
+      hydra/env: default
+      hydra/callbacks: null
+      hydra/job_logging: default
+      hydra/hydra_logging: default
+      hydra/hydra_help: default
+      hydra/help: default
+      hydra/sweeper: basic
+      hydra/launcher: basic
+      hydra/output: default
+  verbose: false

checkpoints/reference_reproduction/_invalid/old_di4c_shape_skip_20260825_095622/duo_di4c_lm1b/.hydra/overrides.yaml

@@ -0,0 +1,40 @@
+- mode=train
+- data.train=lm1b
+- data.valid=lm1b
+- tokenizer.name=/home/ruizzhao/diffusion_baseline/data/raw/huggingface/hub/models--bert-base-uncased/snapshots/86b5e0934494bd15c9632b12f734a8a67f723594
+- data_preprocess.data_cache=/home/ruizzhao/diffusion_baseline/data/training_cache/di4c/lm1b
+- data_preprocess.seq_len=128
+- data_preprocess.group_text=true
+- data_preprocess.remove_text=true
+- data_preprocess.add_bos=false
+- data_preprocess.add_eos=true
+- data_preprocess.legacy_start_end_bos=true
+- model=dit-orig-small
+- model.length=128
+- parameterization=multi-round-sdtt
+- parameterization.checkpoint_path=/home/ruizzhao/diffusion_baseline/checkpoints/reference_reproduction/duo_di4c/lm1b/inputs/teacher_adapted.ckpt
+- parameterization.start_from_hf=false
+- parameterization.grow_dt_every=10000
+- parameterization.orig_num_sampling_steps=1024
+- loader.global_batch_size=2
+- loader.eval_global_batch_size=2
+- loader.batch_size=1
+- loader.eval_batch_size=1
+- trainer.devices=1
+- trainer.num_nodes=1
+- trainer.accumulate_grad_batches=2
+- trainer.max_steps=20000
+- trainer.precision=bf16-mixed
+- optim.lr=3e-05
+- lr_scheduler.num_warmup_steps=2500
+- callbacks.checkpoint_every_n_steps.every_n_train_steps=10000
+- checkpointing.save_dir=/home/ruizzhao/diffusion_baseline/checkpoints/reference_reproduction/duo_di4c/lm1b
+- checkpointing.resume_from_ckpt=true
+- checkpointing.resume_ckpt_path=/home/ruizzhao/diffusion_baseline/checkpoints/reference_reproduction/duo_di4c/lm1b/checkpoints/last.ckpt
+- wandb=null
+- +seed=42
+- is_di4c=true
+- is_teacher_di4c=false
+- T=1024
+- latent_bsize=16
+- round=7

checkpoints/reference_reproduction/_invalid/old_di4c_shape_skip_20260825_095622/duo_di4c_lm1b/config.yaml

@@ -0,0 +1,179 @@
+mode: train
+time_conditioning: false
+T: 1024
+compile: false
+eval_every: 5000000
+is_di4c: true
+is_teacher_di4c: false
+latent_bsize: 16
+alpha_t: sigmoid
+alpha_const: 0.1
+distil_delta: 0.02
+log_unif: false
+round: 7
+data_preprocess:
+  data_cache: /home/ruizzhao/diffusion_baseline/data/training_cache/di4c/lm1b
+  min_seq_len: -1
+  seq_len: 128
+  group_text: true
+  remove_text: true
+  num_seqs: -1
+  add_bos: false
+  add_eos: true
+  legacy_start_end_bos: true
+loader:
+  global_batch_size: 2
+  eval_global_batch_size: 2
+  batch_size: 1
+  eval_batch_size: 1
+  num_workers: ${eval:"len(__import__('os').sched_getaffinity(0))"}
+  pin_memory: true
+  persistent_workers: true
+training:
+  ema: 0.9999
+  antithetic_sampling: true
+  importance_sampling: false
+  sampling_eps: 1.0e-05
+  change_of_variables: false
+tokenizer:
+  name: /home/ruizzhao/diffusion_baseline/data/raw/huggingface/hub/models--bert-base-uncased/snapshots/86b5e0934494bd15c9632b12f734a8a67f723594
+optim:
+  name: adamw
+  weight_decay: 0.0
+  lr: 3.0e-05
+  beta1: 0.9
+  beta2: 0.999
+  eps: 1.0e-08
+trainer:
+  _target_: lightning.Trainer
+  accelerator: cuda
+  num_nodes: 1
+  devices: 1
+  accumulate_grad_batches: 2
+  gradient_clip_val: 1.0
+  precision: bf16-mixed
+  num_sanity_val_steps: 0
+  max_steps: 20000
+  log_every_n_steps: 10
+  limit_train_batches: 1.0
+  limit_val_batches: 1.0
+  val_check_interval: ${eval:${trainer.accumulate_grad_batches} * ${eval_every}}
+  check_val_every_n_epoch: null
+  benchmark: true
+wandb: null
+checkpointing:
+  save_dir: /home/ruizzhao/diffusion_baseline/checkpoints/reference_reproduction/duo_di4c/lm1b
+  resume_from_ckpt: true
+  resume_ckpt_path: /home/ruizzhao/diffusion_baseline/checkpoints/reference_reproduction/duo_di4c/lm1b/checkpoints/last.ckpt
+eval:
+  valid:
+    n_samples: 16
+    num_steps: 128
+  ppl_with_ar:
+    run: true
+    model: gpt2-large
+    batch_size: 8
+  cond_ppl:
+    run: true
+    model: gpt2-large
+    batch_size: 8
+  mauve:
+    run: false
+    num_rounds: 5
+    max_num_tokens: 100
+    batch_size: 8
+    scaling_factor: 5
+  lambada_openai:
+    run: false
+    batch_size: 32
+    num_samples: 32
+    from_ema: false
+    add_bos: true
+    add_eos: true
+callbacks:
+  checkpoint_every_n_steps:
+    _target_: lightning.pytorch.callbacks.ModelCheckpoint
+    save_top_k: -1
+    save_last: true
+    dirpath: ${checkpointing.save_dir}/checkpoints
+    verbose: true
+    auto_insert_metric_name: false
+    every_n_train_steps: 10000
+  learning_rate_monitor:
+    _target_: lightning.pytorch.callbacks.LearningRateMonitor
+    logging_interval: step
+  text_throughput:
+    _target_: callbacks.text_throughput.TextThroughputMonitor
+  grad_norm:
+    _target_: callbacks.gradient_norm.GradNormCallback
+data:
+  train: lm1b
+  valid: lm1b
+  test: openwebtext-test
+  tokenizer_name_or_path: gpt2
+  wrap: true
+  streaming: false
+model:
+  name: small
+  type: ddit-orig
+  hidden_size: 768
+  cond_dim: 128
+  length: 128
+  n_blocks: 12
+  n_heads: 12
+  scale_by_sigma: true
+  dropout: 0.1
+  tie_word_embeddings: false
+strategy:
+  _target_: lightning.pytorch.strategies.DDPStrategy
+  find_unused_parameters: true
+noise:
+  type: loglinear
+  sigma_min: 0.0001
+  sigma_max: 20
+lr_scheduler:
+  _target_: transformers.get_constant_schedule_with_warmup
+  num_warmup_steps: 2500
+parameterization:
+  name: multi-round-sdtt
+  log_loss_buckets: -1
+  checkpoint_path: /home/ruizzhao/diffusion_baseline/checkpoints/reference_reproduction/duo_di4c/lm1b/inputs/teacher_adapted.ckpt
+  start_from_hf: false
+  distill_mode: kl-bwd
+  num_distill_steps: 2
+  min_num_sampling_steps: 8
+  grow_dt_every: 10000
+  orig_num_sampling_steps: 1024
+  sampling_mode: ancestral
+  loss_precision: null
+  reset_optimizer_on_growth: false
+  use_ema_on_growth: false
+  sampling:
+    uncond:
+      run: true
+      num_samples: 1024
+      batch_size: 16
+      from_ema: false
+      num_steps: 16
+      seq_len: 1024
+      sampler: ancestral
+      cache_preds: false
+      add_bos: false
+      add_eos: false
+    cond_prefix:
+      run: true
+      num_samples: 256
+      batch_size: 16
+      from_ema: false
+      dataset: webtext
+      seq_len: 100
+      prefix_len: 50
+      num_cont_per_prefix: 5
+      min_seq_len: 1024
+      num_steps: 128
+      sampler: ancestral
+      cache_preds: false
+      add_bos: false
+      add_eos: false
+      calc_sbleu: true
+seed: 42

checkpoints/reference_reproduction/_invalid/old_di4c_shape_skip_20260825_095622/duo_di4c_lm1b/logs_di4c_train.txt

@@ -0,0 +1,53 @@
+2026-08-20 17:43:21.308 | INFO     | lightning_utilities.core.rank_zero:wrapped_fn:42 - Starting training
+2026-08-20 17:43:21.309 | INFO     | lightning_utilities.core.rank_zero:wrapped_fn:42 - Training starting from scratch (no checkpoint to reload)
+2026-08-20 17:43:22.879 | WARNING  | lightning_utilities.core.rank_zero:wrapped_fn:42 - Skipping HF student tensors with incompatible shapes: [('backbone.output_layer.linear.bias', (50258,), (30523,)), ('backbone.output_layer.linear.weight', (50258, 768), (30523, 768)), ('backbone.vocab_embed.embedding', (50258, 768), (30523, 768))]
+2026-08-20 17:43:30.787 | WARNING  | lightning_utilities.core.rank_zero:wrapped_fn:42 - Using Legacy dataloader!!!
+2026-08-20 17:43:30.791 | INFO     | lightning_utilities.core.rank_zero:wrapped_fn:42 - Loading data from: /home/ruizzhao/diffusion_baseline/data/training_cache/di4c/lm1b/lm1b_train_bs128_wrapped.dat
+2026-08-20 17:43:30.830 | WARNING  | lightning_utilities.core.rank_zero:wrapped_fn:42 - Using Legacy dataloader!!!
+2026-08-20 17:43:30.830 | WARNING  | lightning_utilities.core.rank_zero:wrapped_fn:42 - Using LM1B test set as validation set
+2026-08-20 17:43:30.830 | INFO     | lightning_utilities.core.rank_zero:wrapped_fn:42 - Loading data from: /home/ruizzhao/diffusion_baseline/data/training_cache/di4c/lm1b/lm1b_test_bs128_wrapped.dat
+2026-08-20 17:43:30.832 | INFO     | lightning_utilities.core.rank_zero:wrapped_fn:42 - Printing train batch.
+2026-08-20 17:43:30.833 | INFO     | lightning_utilities.core.rank_zero:wrapped_fn:42 - Batch input_ids.shape: torch.Size([1, 128])
+2026-08-20 17:43:30.833 | INFO     | lightning_utilities.core.rank_zero:wrapped_fn:42 - First 64 tokens: [CLS] while athletes in different professions dealt with doping scandals and other controversies, woods continued to do what he did best : dominate the field of professional golf and rake in endorsements. [SEP] the minutes could shed light on an internal debate, which has been evident in fed officials'recent speeches, over when to consider raising rates.
+2026-08-20 17:43:30.834 | INFO     | lightning_utilities.core.rank_zero:wrapped_fn:42 - ids: tensor([  101,  2096,  7576,  1999,  2367, 22797,  9411,  2007, 23799, 29609,
+         1998,  2060, 25962,  1010,  5249,  2506,  2000,  2079,  2054,  2002,
+         2106,  2190,  1024, 16083,  1996,  2492,  1997,  2658,  5439,  1998,
+        26008,  1999, 20380,  2015,  1012,   102,  1996,  2781,  2071,  8328,
+         2422,  2006,  2019,  4722,  5981,  1010,  2029,  2038,  2042, 10358,
+         1999,  7349,  4584,  1005,  3522, 13867,  1010,  2058,  2043,  2000,
+         5136,  6274,  6165,  1012])
+2026-08-20 17:43:30.834 | INFO     | lightning_utilities.core.rank_zero:wrapped_fn:42 - Last 64 tokens: [SEP] but metro's definition of " urbanite " is not accurate, says matthew gandy, director of the urban laboratory at university college london. [SEP] goodrich petroleum is an independent oil and gas exploration and production company listed on the new york stock exchange. [SEP] the support from spears is to end on november 15 [SEP]
+2026-08-20 17:43:30.834 | INFO     | lightning_utilities.core.rank_zero:wrapped_fn:42 - ids: tensor([  102,  2021,  6005,  1005,  1055,  6210,  1997,  1000,  3923,  4221,
+         1000,  2003,  2025,  8321,  1010,  2758,  5487, 25957,  5149,  1010,
+         2472,  1997,  1996,  3923,  5911,  2012,  2118,  2267,  2414,  1012,
+          102,  2204, 13149, 11540,  2003,  2019,  2981,  3514,  1998,  3806,
+         8993,  1998,  2537,  2194,  3205,  2006,  1996,  2047,  2259,  4518,
+         3863,  1012,   102,  1996,  2490,  2013, 13957,  2003,  2000,  2203,
+         2006,  2281,  2321,   102])
+2026-08-20 17:43:30.834 | INFO     | lightning_utilities.core.rank_zero:wrapped_fn:42 - ==================================================
+2026-08-20 17:43:30.834 | INFO     | lightning_utilities.core.rank_zero:wrapped_fn:42 - Printing valid batch.
+2026-08-20 17:43:30.835 | INFO     | lightning_utilities.core.rank_zero:wrapped_fn:42 - Batch input_ids.shape: torch.Size([1, 128])
+2026-08-20 17:43:30.835 | INFO     | lightning_utilities.core.rank_zero:wrapped_fn:42 - First 64 tokens: [CLS] in medieval times, the townspeople from norcia were so practised at butchering pigs that they gained an in - depth knowledge of anatomy and were allowed to practise as surgeons along with members of the clergy. [SEP] the satellite weighs some 660 pounds, the israeli haaretz newspaper reported, citing unnamed company officials.
+2026-08-20 17:43:30.835 | INFO     | lightning_utilities.core.rank_zero:wrapped_fn:42 - ids: tensor([  101,  1999,  5781,  2335,  1010,  1996, 27938,  2013,  4496,  7405,
+         2020,  2061, 20439,  2012, 14998,  2075, 14695,  2008,  2027,  4227,
+         2019,  1999,  1011,  5995,  3716,  1997, 13336,  1998,  2020,  3039,
+         2000, 10975, 18908,  5562,  2004, 16804,  2247,  2007,  2372,  1997,
+         1996, 11646,  1012,   102,  1996,  5871, 21094,  2070, 20982,  7038,
+         1010,  1996,  5611,  5292, 12069,  5753,  3780,  2988,  1010,  8951,
+        13294,  2194,  4584,  1012])
+2026-08-20 17:43:30.835 | INFO     | lightning_utilities.core.rank_zero:wrapped_fn:42 - Last 64 tokens: [SEP] the situation has been further complicated by growing indications that the us treasury is preparing to make funds available for recapitalsing us banks. [SEP] as north korea's close neighbor, traditional ally and main provider of economic aid, china is widely believed to hold the key to solving the north korea conundrum. [SEP]
+2026-08-20 17:43:30.835 | INFO     | lightning_utilities.core.rank_zero:wrapped_fn:42 - ids: tensor([  102,  1996,  3663,  2038,  2042,  2582,  8552,  2011,  3652, 24936,
+         2008,  1996,  2149,  9837,  2003,  8225,  2000,  2191,  5029,  2800,
+         2005, 28667,  9331, 18400,  7741,  2149,  5085,  1012,   102,  2004,
+         2167,  4420,  1005,  1055,  2485, 11429,  1010,  3151,  9698,  1998,
+         2364, 10802,  1997,  3171,  4681,  1010,  2859,  2003,  4235,  3373,
+         2000,  2907,  1996,  3145,  2000, 13729,  1996,  2167,  4420,  9530,
+         8630,  6824,  1012,   102])
+2026-08-20 17:43:30.835 | INFO     | lightning_utilities.core.rank_zero:wrapped_fn:42 - ==================================================
+2026-08-20 17:43:30.835 | INFO     | lightning_utilities.core.rank_zero:wrapped_fn:42 - Train set length: 7102854
+2026-08-20 17:43:30.835 | INFO     | lightning_utilities.core.rank_zero:wrapped_fn:42 - Valid set length: 71979
+2026-08-20 17:43:30.853 | WARNING  | lightning_utilities.core.rank_zero:wrapped_fn:42 - Using Legacy dataloader!!!
+2026-08-20 17:43:30.854 | INFO     | lightning_utilities.core.rank_zero:wrapped_fn:42 - Loading data from: /home/ruizzhao/diffusion_baseline/data/training_cache/di4c/lm1b/lm1b_train_bs128_wrapped.dat
+2026-08-20 17:43:30.936 | WARNING  | lightning_utilities.core.rank_zero:wrapped_fn:42 - Using Legacy dataloader!!!
+2026-08-20 17:43:30.936 | WARNING  | lightning_utilities.core.rank_zero:wrapped_fn:42 - Using LM1B test set as validation set
+2026-08-20 17:43:30.936 | INFO     | lightning_utilities.core.rank_zero:wrapped_fn:42 - Loading data from: /home/ruizzhao/diffusion_baseline/data/training_cache/di4c/lm1b/lm1b_test_bs128_wrapped.dat

checkpoints/reference_reproduction/_invalid/old_di4c_shape_skip_20260825_095622/mdlm_di4c_lm1b/.hydra/config.yaml

@@ -0,0 +1,179 @@
+mode: train
+time_conditioning: false
+T: 1024
+compile: false
+eval_every: 5000000
+is_di4c: true
+is_teacher_di4c: false
+latent_bsize: 16
+alpha_t: sigmoid
+alpha_const: 0.1
+distil_delta: 0.02
+log_unif: false
+round: 7
+data_preprocess:
+  data_cache: /home/ruizzhao/diffusion_baseline/data/training_cache/di4c/lm1b
+  min_seq_len: -1
+  seq_len: 128
+  group_text: true
+  remove_text: true
+  num_seqs: -1
+  add_bos: false
+  add_eos: true
+  legacy_start_end_bos: true
+loader:
+  global_batch_size: 2
+  eval_global_batch_size: 2
+  batch_size: 1
+  eval_batch_size: 1
+  num_workers: ${eval:"len(__import__('os').sched_getaffinity(0))"}
+  pin_memory: true
+  persistent_workers: true
+training:
+  ema: 0.9999
+  antithetic_sampling: true
+  importance_sampling: false
+  sampling_eps: 1.0e-05
+  change_of_variables: false
+tokenizer:
+  name: /home/ruizzhao/diffusion_baseline/data/raw/huggingface/hub/models--bert-base-uncased/snapshots/86b5e0934494bd15c9632b12f734a8a67f723594
+optim:
+  name: adamw
+  weight_decay: 0.0
+  lr: 3.0e-05
+  beta1: 0.9
+  beta2: 0.999
+  eps: 1.0e-08
+trainer:
+  _target_: lightning.Trainer
+  accelerator: cuda
+  num_nodes: 1
+  devices: 1
+  accumulate_grad_batches: 2
+  gradient_clip_val: 1.0
+  precision: bf16-mixed
+  num_sanity_val_steps: 0
+  max_steps: 20000
+  log_every_n_steps: 10
+  limit_train_batches: 1.0
+  limit_val_batches: 1.0
+  val_check_interval: ${eval:${trainer.accumulate_grad_batches} * ${eval_every}}
+  check_val_every_n_epoch: null
+  benchmark: true
+wandb: null
+checkpointing:
+  save_dir: /home/ruizzhao/diffusion_baseline/checkpoints/reference_reproduction/mdlm_di4c/lm1b
+  resume_from_ckpt: true
+  resume_ckpt_path: /home/ruizzhao/diffusion_baseline/checkpoints/reference_reproduction/mdlm_di4c/lm1b/checkpoints/last.ckpt
+eval:
+  valid:
+    n_samples: 16
+    num_steps: 128
+  ppl_with_ar:
+    run: true
+    model: gpt2-large
+    batch_size: 8
+  cond_ppl:
+    run: true
+    model: gpt2-large
+    batch_size: 8
+  mauve:
+    run: false
+    num_rounds: 5
+    max_num_tokens: 100
+    batch_size: 8
+    scaling_factor: 5
+  lambada_openai:
+    run: false
+    batch_size: 32
+    num_samples: 32
+    from_ema: false
+    add_bos: true
+    add_eos: true
+callbacks:
+  checkpoint_every_n_steps:
+    _target_: lightning.pytorch.callbacks.ModelCheckpoint
+    save_top_k: -1
+    save_last: true
+    dirpath: ${checkpointing.save_dir}/checkpoints
+    verbose: true
+    auto_insert_metric_name: false
+    every_n_train_steps: 10000
+  learning_rate_monitor:
+    _target_: lightning.pytorch.callbacks.LearningRateMonitor
+    logging_interval: step
+  text_throughput:
+    _target_: callbacks.text_throughput.TextThroughputMonitor
+  grad_norm:
+    _target_: callbacks.gradient_norm.GradNormCallback
+data:
+  train: lm1b
+  valid: lm1b
+  test: openwebtext-test
+  tokenizer_name_or_path: gpt2
+  wrap: true
+  streaming: false
+model:
+  name: small
+  type: ddit-orig
+  hidden_size: 768
+  cond_dim: 128
+  length: 128
+  n_blocks: 12
+  n_heads: 12
+  scale_by_sigma: true
+  dropout: 0.1
+  tie_word_embeddings: false
+strategy:
+  _target_: lightning.pytorch.strategies.DDPStrategy
+  find_unused_parameters: true
+noise:
+  type: loglinear
+  sigma_min: 0.0001
+  sigma_max: 20
+lr_scheduler:
+  _target_: transformers.get_constant_schedule_with_warmup
+  num_warmup_steps: 2500
+parameterization:
+  name: multi-round-sdtt
+  log_loss_buckets: -1
+  checkpoint_path: /home/ruizzhao/diffusion_baseline/checkpoints/reference_reproduction/mdlm_di4c/lm1b/inputs/teacher_adapted.ckpt
+  start_from_hf: false
+  distill_mode: kl-bwd
+  num_distill_steps: 2
+  min_num_sampling_steps: 8
+  grow_dt_every: 10000
+  orig_num_sampling_steps: 1024
+  sampling_mode: ancestral
+  loss_precision: null
+  reset_optimizer_on_growth: false
+  use_ema_on_growth: false
+  sampling:
+    uncond:
+      run: true
+      num_samples: 1024
+      batch_size: 16
+      from_ema: false
+      num_steps: 16
+      seq_len: 1024
+      sampler: ancestral
+      cache_preds: false
+      add_bos: false
+      add_eos: false
+    cond_prefix:
+      run: true
+      num_samples: 256
+      batch_size: 16
+      from_ema: false
+      dataset: webtext
+      seq_len: 100
+      prefix_len: 50
+      num_cont_per_prefix: 5
+      min_seq_len: 1024
+      num_steps: 128
+      sampler: ancestral
+      cache_preds: false
+      add_bos: false
+      add_eos: false
+      calc_sbleu: true
+seed: 42

checkpoints/reference_reproduction/_invalid/old_di4c_shape_skip_20260825_095622/mdlm_di4c_lm1b/.hydra/hydra.yaml

@@ -0,0 +1,202 @@
+hydra:
+  run:
+    dir: /home/ruizzhao/diffusion_baseline/checkpoints/reference_reproduction/mdlm_di4c/lm1b
+  sweep:
+    dir: multirun/${now:%Y-%m-%d}/${now:%H-%M-%S}
+    subdir: ${hydra.job.num}
+  launcher:
+    _target_: hydra._internal.core_plugins.basic_launcher.BasicLauncher
+  sweeper:
+    _target_: hydra._internal.core_plugins.basic_sweeper.BasicSweeper
+    max_batch_size: null
+    params: null
+  help:
+    app_name: ${hydra.job.name}
+    header: '${hydra.help.app_name} is powered by Hydra.
+
+      '
+    footer: 'Powered by Hydra (https://hydra.cc)
+
+      Use --hydra-help to view Hydra specific help
+
+      '
+    template: '${hydra.help.header}
+
+      == Configuration groups ==
+
+      Compose your configuration from those groups (group=option)
+
+
+      $APP_CONFIG_GROUPS
+
+
+      == Config ==
+
+      Override anything in the config (foo.bar=value)
+
+
+      $CONFIG
+
+
+      ${hydra.help.footer}
+
+      '
+  hydra_help:
+    template: 'Hydra (${hydra.runtime.version})
+
+      See https://hydra.cc for more info.
+
+
+      == Flags ==
+
+      $FLAGS_HELP
+
+
+      == Configuration groups ==
+
+      Compose your configuration from those groups (For example, append hydra/job_logging=disabled
+      to command line)
+
+
+      $HYDRA_CONFIG_GROUPS
+
+
+      Use ''--cfg hydra'' to Show the Hydra config.
+
+      '
+    hydra_help: ???
+  hydra_logging:
+    version: 1
+    formatters:
+      simple:
+        format: '[%(asctime)s][HYDRA] %(message)s'
+    handlers:
+      console:
+        class: logging.StreamHandler
+        formatter: simple
+        stream: ext://sys.stdout
+    root:
+      level: INFO
+      handlers:
+      - console
+    loggers:
+      logging_example:
+        level: DEBUG
+    disable_existing_loggers: false
+  job_logging:
+    version: 1
+    formatters:
+      simple:
+        format: '[%(asctime)s][%(name)s][%(levelname)s] - %(message)s'
+    handlers:
+      console:
+        class: logging.StreamHandler
+        formatter: simple
+        stream: ext://sys.stdout
+      file:
+        class: logging.FileHandler
+        formatter: simple
+        filename: ${hydra.runtime.output_dir}/${hydra.job.name}.log
+    root:
+      level: INFO
+      handlers:
+      - console
+      - file
+    disable_existing_loggers: false
+  env: {}
+  mode: RUN
+  searchpath: []
+  callbacks: {}
+  output_subdir: .hydra
+  overrides:
+    hydra:
+    - hydra.run.dir=/home/ruizzhao/diffusion_baseline/checkpoints/reference_reproduction/mdlm_di4c/lm1b
+    - hydra.job.chdir=true
+    - hydra.mode=RUN
+    task:
+    - mode=train
+    - data.train=lm1b
+    - data.valid=lm1b
+    - tokenizer.name=/home/ruizzhao/diffusion_baseline/data/raw/huggingface/hub/models--bert-base-uncased/snapshots/86b5e0934494bd15c9632b12f734a8a67f723594
+    - data_preprocess.data_cache=/home/ruizzhao/diffusion_baseline/data/training_cache/di4c/lm1b
+    - data_preprocess.seq_len=128
+    - data_preprocess.group_text=true
+    - data_preprocess.remove_text=true
+    - data_preprocess.add_bos=false
+    - data_preprocess.add_eos=true
+    - data_preprocess.legacy_start_end_bos=true
+    - model=dit-orig-small
+    - model.length=128
+    - parameterization=multi-round-sdtt
+    - parameterization.checkpoint_path=/home/ruizzhao/diffusion_baseline/checkpoints/reference_reproduction/mdlm_di4c/lm1b/inputs/teacher_adapted.ckpt
+    - parameterization.start_from_hf=false
+    - parameterization.grow_dt_every=10000
+    - parameterization.orig_num_sampling_steps=1024
+    - loader.global_batch_size=2
+    - loader.eval_global_batch_size=2
+    - loader.batch_size=1
+    - loader.eval_batch_size=1
+    - trainer.devices=1
+    - trainer.num_nodes=1
+    - trainer.accumulate_grad_batches=2
+    - trainer.max_steps=20000
+    - trainer.precision=bf16-mixed
+    - optim.lr=3e-05
+    - lr_scheduler.num_warmup_steps=2500
+    - callbacks.checkpoint_every_n_steps.every_n_train_steps=10000
+    - checkpointing.save_dir=/home/ruizzhao/diffusion_baseline/checkpoints/reference_reproduction/mdlm_di4c/lm1b
+    - checkpointing.resume_from_ckpt=true
+    - checkpointing.resume_ckpt_path=/home/ruizzhao/diffusion_baseline/checkpoints/reference_reproduction/mdlm_di4c/lm1b/checkpoints/last.ckpt
+    - wandb=null
+    - +seed=42
+    - is_di4c=true
+    - is_teacher_di4c=false
+    - T=1024
+    - latent_bsize=16
+    - round=7
+  job:
+    name: main
+    chdir: true
+    override_dirname: +seed=42,T=1024,callbacks.checkpoint_every_n_steps.every_n_train_steps=10000,checkpointing.resume_ckpt_path=/home/ruizzhao/diffusion_baseline/checkpoints/reference_reproduction/mdlm_di4c/lm1b/checkpoints/last.ckpt,checkpointing.resume_from_ckpt=true,checkpointing.save_dir=/home/ruizzhao/diffusion_baseline/checkpoints/reference_reproduction/mdlm_di4c/lm1b,data.train=lm1b,data.valid=lm1b,data_preprocess.add_bos=false,data_preprocess.add_eos=true,data_preprocess.data_cache=/home/ruizzhao/diffusion_baseline/data/training_cache/di4c/lm1b,data_preprocess.group_text=true,data_preprocess.legacy_start_end_bos=true,data_preprocess.remove_text=true,data_preprocess.seq_len=128,is_di4c=true,is_teacher_di4c=false,latent_bsize=16,loader.batch_size=1,loader.eval_batch_size=1,loader.eval_global_batch_size=2,loader.global_batch_size=2,lr_scheduler.num_warmup_steps=2500,mode=train,model.length=128,model=dit-orig-small,optim.lr=3e-05,parameterization.checkpoint_path=/home/ruizzhao/diffusion_baseline/checkpoints/reference_reproduction/mdlm_di4c/lm1b/inputs/teacher_adapted.ckpt,parameterization.grow_dt_every=10000,parameterization.orig_num_sampling_steps=1024,parameterization.start_from_hf=false,parameterization=multi-round-sdtt,round=7,tokenizer.name=/home/ruizzhao/diffusion_baseline/data/raw/huggingface/hub/models--bert-base-uncased/snapshots/86b5e0934494bd15c9632b12f734a8a67f723594,trainer.accumulate_grad_batches=2,trainer.devices=1,trainer.max_steps=20000,trainer.num_nodes=1,trainer.precision=bf16-mixed,wandb=null
+    id: ???
+    num: ???
+    config_name: config
+    env_set: {}
+    env_copy: []
+    config:
+      override_dirname:
+        kv_sep: '='
+        item_sep: ','
+        exclude_keys: []
+  runtime:
+    version: 1.3.2
+    version_base: '1.3'
+    cwd: /home/ruizzhao/diffusion_baseline/upstreams/di4c
+    config_sources:
+    - path: hydra.conf
+      schema: pkg
+      provider: hydra
+    - path: /home/ruizzhao/diffusion_baseline/upstreams/di4c/sdtt/src/sdtt/configs
+      schema: file
+      provider: main
+    - path: ''
+      schema: structured
+      provider: schema
+    output_dir: /home/ruizzhao/diffusion_baseline/checkpoints/reference_reproduction/mdlm_di4c/lm1b
+    choices:
+      parameterization: multi-round-sdtt
+      lr_scheduler: constant_warmup
+      noise: loglinear
+      strategy: ddp
+      model: dit-orig-small
+      data: openwebtext-split
+      hydra/env: default
+      hydra/callbacks: null
+      hydra/job_logging: default
+      hydra/hydra_logging: default
+      hydra/hydra_help: default
+      hydra/help: default
+      hydra/sweeper: basic
+      hydra/launcher: basic
+      hydra/output: default
+  verbose: false

checkpoints/reference_reproduction/_invalid/old_di4c_shape_skip_20260825_095622/mdlm_di4c_lm1b/.hydra/overrides.yaml

@@ -0,0 +1,40 @@
+- mode=train
+- data.train=lm1b
+- data.valid=lm1b
+- tokenizer.name=/home/ruizzhao/diffusion_baseline/data/raw/huggingface/hub/models--bert-base-uncased/snapshots/86b5e0934494bd15c9632b12f734a8a67f723594
+- data_preprocess.data_cache=/home/ruizzhao/diffusion_baseline/data/training_cache/di4c/lm1b
+- data_preprocess.seq_len=128
+- data_preprocess.group_text=true
+- data_preprocess.remove_text=true
+- data_preprocess.add_bos=false
+- data_preprocess.add_eos=true
+- data_preprocess.legacy_start_end_bos=true
+- model=dit-orig-small
+- model.length=128
+- parameterization=multi-round-sdtt
+- parameterization.checkpoint_path=/home/ruizzhao/diffusion_baseline/checkpoints/reference_reproduction/mdlm_di4c/lm1b/inputs/teacher_adapted.ckpt
+- parameterization.start_from_hf=false
+- parameterization.grow_dt_every=10000
+- parameterization.orig_num_sampling_steps=1024
+- loader.global_batch_size=2
+- loader.eval_global_batch_size=2
+- loader.batch_size=1
+- loader.eval_batch_size=1
+- trainer.devices=1
+- trainer.num_nodes=1
+- trainer.accumulate_grad_batches=2
+- trainer.max_steps=20000
+- trainer.precision=bf16-mixed
+- optim.lr=3e-05
+- lr_scheduler.num_warmup_steps=2500
+- callbacks.checkpoint_every_n_steps.every_n_train_steps=10000
+- checkpointing.save_dir=/home/ruizzhao/diffusion_baseline/checkpoints/reference_reproduction/mdlm_di4c/lm1b
+- checkpointing.resume_from_ckpt=true
+- checkpointing.resume_ckpt_path=/home/ruizzhao/diffusion_baseline/checkpoints/reference_reproduction/mdlm_di4c/lm1b/checkpoints/last.ckpt
+- wandb=null
+- +seed=42
+- is_di4c=true
+- is_teacher_di4c=false
+- T=1024
+- latent_bsize=16
+- round=7

checkpoints/reference_reproduction/_invalid/old_di4c_shape_skip_20260825_095622/mdlm_di4c_lm1b/config.yaml

@@ -0,0 +1,179 @@
+mode: train
+time_conditioning: false
+T: 1024
+compile: false
+eval_every: 5000000
+is_di4c: true
+is_teacher_di4c: false
+latent_bsize: 16
+alpha_t: sigmoid
+alpha_const: 0.1
+distil_delta: 0.02
+log_unif: false
+round: 7
+data_preprocess:
+  data_cache: /home/ruizzhao/diffusion_baseline/data/training_cache/di4c/lm1b
+  min_seq_len: -1
+  seq_len: 128
+  group_text: true
+  remove_text: true
+  num_seqs: -1
+  add_bos: false
+  add_eos: true
+  legacy_start_end_bos: true
+loader:
+  global_batch_size: 2
+  eval_global_batch_size: 2
+  batch_size: 1
+  eval_batch_size: 1
+  num_workers: ${eval:"len(__import__('os').sched_getaffinity(0))"}
+  pin_memory: true
+  persistent_workers: true
+training:
+  ema: 0.9999
+  antithetic_sampling: true
+  importance_sampling: false
+  sampling_eps: 1.0e-05
+  change_of_variables: false
+tokenizer:
+  name: /home/ruizzhao/diffusion_baseline/data/raw/huggingface/hub/models--bert-base-uncased/snapshots/86b5e0934494bd15c9632b12f734a8a67f723594
+optim:
+  name: adamw
+  weight_decay: 0.0
+  lr: 3.0e-05
+  beta1: 0.9
+  beta2: 0.999
+  eps: 1.0e-08
+trainer:
+  _target_: lightning.Trainer
+  accelerator: cuda
+  num_nodes: 1
+  devices: 1
+  accumulate_grad_batches: 2
+  gradient_clip_val: 1.0
+  precision: bf16-mixed
+  num_sanity_val_steps: 0
+  max_steps: 20000
+  log_every_n_steps: 10
+  limit_train_batches: 1.0
+  limit_val_batches: 1.0
+  val_check_interval: ${eval:${trainer.accumulate_grad_batches} * ${eval_every}}
+  check_val_every_n_epoch: null
+  benchmark: true
+wandb: null
+checkpointing:
+  save_dir: /home/ruizzhao/diffusion_baseline/checkpoints/reference_reproduction/mdlm_di4c/lm1b
+  resume_from_ckpt: true
+  resume_ckpt_path: /home/ruizzhao/diffusion_baseline/checkpoints/reference_reproduction/mdlm_di4c/lm1b/checkpoints/last.ckpt
+eval:
+  valid:
+    n_samples: 16
+    num_steps: 128
+  ppl_with_ar:
+    run: true
+    model: gpt2-large
+    batch_size: 8
+  cond_ppl:
+    run: true
+    model: gpt2-large
+    batch_size: 8
+  mauve:
+    run: false
+    num_rounds: 5
+    max_num_tokens: 100
+    batch_size: 8
+    scaling_factor: 5
+  lambada_openai:
+    run: false
+    batch_size: 32
+    num_samples: 32
+    from_ema: false
+    add_bos: true
+    add_eos: true
+callbacks:
+  checkpoint_every_n_steps:
+    _target_: lightning.pytorch.callbacks.ModelCheckpoint
+    save_top_k: -1
+    save_last: true
+    dirpath: ${checkpointing.save_dir}/checkpoints
+    verbose: true
+    auto_insert_metric_name: false
+    every_n_train_steps: 10000
+  learning_rate_monitor:
+    _target_: lightning.pytorch.callbacks.LearningRateMonitor
+    logging_interval: step
+  text_throughput:
+    _target_: callbacks.text_throughput.TextThroughputMonitor
+  grad_norm:
+    _target_: callbacks.gradient_norm.GradNormCallback
+data:
+  train: lm1b
+  valid: lm1b
+  test: openwebtext-test
+  tokenizer_name_or_path: gpt2
+  wrap: true
+  streaming: false
+model:
+  name: small
+  type: ddit-orig
+  hidden_size: 768
+  cond_dim: 128
+  length: 128
+  n_blocks: 12
+  n_heads: 12
+  scale_by_sigma: true
+  dropout: 0.1
+  tie_word_embeddings: false
+strategy:
+  _target_: lightning.pytorch.strategies.DDPStrategy
+  find_unused_parameters: true
+noise:
+  type: loglinear
+  sigma_min: 0.0001
+  sigma_max: 20
+lr_scheduler:
+  _target_: transformers.get_constant_schedule_with_warmup
+  num_warmup_steps: 2500
+parameterization:
+  name: multi-round-sdtt
+  log_loss_buckets: -1
+  checkpoint_path: /home/ruizzhao/diffusion_baseline/checkpoints/reference_reproduction/mdlm_di4c/lm1b/inputs/teacher_adapted.ckpt
+  start_from_hf: false
+  distill_mode: kl-bwd
+  num_distill_steps: 2
+  min_num_sampling_steps: 8
+  grow_dt_every: 10000
+  orig_num_sampling_steps: 1024
+  sampling_mode: ancestral
+  loss_precision: null
+  reset_optimizer_on_growth: false
+  use_ema_on_growth: false
+  sampling:
+    uncond:
+      run: true
+      num_samples: 1024
+      batch_size: 16
+      from_ema: false
+      num_steps: 16
+      seq_len: 1024
+      sampler: ancestral
+      cache_preds: false
+      add_bos: false
+      add_eos: false
+    cond_prefix:
+      run: true
+      num_samples: 256
+      batch_size: 16
+      from_ema: false
+      dataset: webtext
+      seq_len: 100
+      prefix_len: 50
+      num_cont_per_prefix: 5
+      min_seq_len: 1024
+      num_steps: 128
+      sampler: ancestral
+      cache_preds: false
+      add_bos: false
+      add_eos: false
+      calc_sbleu: true
+seed: 42

checkpoints/reference_reproduction/_invalid/old_di4c_shape_skip_20260825_095622/mdlm_di4c_lm1b/inputs/teacher_adapted.provenance.json

@@ -0,0 +1,16 @@
+{
+  "hidden_size": 768,
+  "output_path": "/home/ruizzhao/diffusion_baseline/checkpoints/reference_reproduction/mdlm_di4c/lm1b/inputs/teacher_adapted.ckpt",
+  "output_sha256": "4e97e6db6a0fbe5ab0f7a86d1e6990664303755ca6eb60a0c7976ee88e78a3ec",
+  "schema_version": 1,
+  "source_mask_index": 103,
+  "source_path": "/home/ruizzhao/diffusion_baseline/checkpoints/reference_reproduction/flm_baselines/lm1b/lm1b_MDLM_.ckpt",
+  "source_sha256": "d42739e5a99bbd7c0f604dff7fe4c9ab12348b475e30f1d807e70ef32249e514",
+  "source_vocab_size": 30522,
+  "state_keys_sha256": "8212d507b6931a59375f8127f0ccb9e51c24da74fe05e974648d92c5612fb803",
+  "target_mask_index": 30522,
+  "target_vocab_size": 30523,
+  "teacher_family": "masked_mdlm",
+  "transformation": "masked_to_absorbing",
+  "transformation_sha256": "0d4b7caaa7c4c2de75b488026be396440bb4ea6c4dd11dd3edfbfb83a3364711"
+}

checkpoints/reference_reproduction/_invalid/old_di4c_shape_skip_20260825_095622/mdlm_di4c_lm1b/logs_di4c_train.txt

@@ -0,0 +1,53 @@
+2026-08-20 19:11:57.885 | INFO     | lightning_utilities.core.rank_zero:wrapped_fn:42 - Starting training
+2026-08-20 19:11:57.885 | INFO     | lightning_utilities.core.rank_zero:wrapped_fn:42 - Training starting from scratch (no checkpoint to reload)
+2026-08-20 19:12:00.016 | WARNING  | lightning_utilities.core.rank_zero:wrapped_fn:42 - Skipping HF student tensors with incompatible shapes: [('backbone.output_layer.linear.bias', (50258,), (30523,)), ('backbone.output_layer.linear.weight', (50258, 768), (30523, 768)), ('backbone.vocab_embed.embedding', (50258, 768), (30523, 768))]
+2026-08-20 19:12:06.762 | WARNING  | lightning_utilities.core.rank_zero:wrapped_fn:42 - Using Legacy dataloader!!!
+2026-08-20 19:12:06.764 | INFO     | lightning_utilities.core.rank_zero:wrapped_fn:42 - Loading data from: /home/ruizzhao/diffusion_baseline/data/training_cache/di4c/lm1b/lm1b_train_bs128_wrapped.dat
+2026-08-20 19:12:06.815 | WARNING  | lightning_utilities.core.rank_zero:wrapped_fn:42 - Using Legacy dataloader!!!
+2026-08-20 19:12:06.815 | WARNING  | lightning_utilities.core.rank_zero:wrapped_fn:42 - Using LM1B test set as validation set
+2026-08-20 19:12:06.815 | INFO     | lightning_utilities.core.rank_zero:wrapped_fn:42 - Loading data from: /home/ruizzhao/diffusion_baseline/data/training_cache/di4c/lm1b/lm1b_test_bs128_wrapped.dat
+2026-08-20 19:12:06.818 | INFO     | lightning_utilities.core.rank_zero:wrapped_fn:42 - Printing train batch.
+2026-08-20 19:12:06.819 | INFO     | lightning_utilities.core.rank_zero:wrapped_fn:42 - Batch input_ids.shape: torch.Size([1, 128])
+2026-08-20 19:12:06.819 | INFO     | lightning_utilities.core.rank_zero:wrapped_fn:42 - First 64 tokens: [CLS] while athletes in different professions dealt with doping scandals and other controversies, woods continued to do what he did best : dominate the field of professional golf and rake in endorsements. [SEP] the minutes could shed light on an internal debate, which has been evident in fed officials'recent speeches, over when to consider raising rates.
+2026-08-20 19:12:06.820 | INFO     | lightning_utilities.core.rank_zero:wrapped_fn:42 - ids: tensor([  101,  2096,  7576,  1999,  2367, 22797,  9411,  2007, 23799, 29609,
+         1998,  2060, 25962,  1010,  5249,  2506,  2000,  2079,  2054,  2002,
+         2106,  2190,  1024, 16083,  1996,  2492,  1997,  2658,  5439,  1998,
+        26008,  1999, 20380,  2015,  1012,   102,  1996,  2781,  2071,  8328,
+         2422,  2006,  2019,  4722,  5981,  1010,  2029,  2038,  2042, 10358,
+         1999,  7349,  4584,  1005,  3522, 13867,  1010,  2058,  2043,  2000,
+         5136,  6274,  6165,  1012])
+2026-08-20 19:12:06.820 | INFO     | lightning_utilities.core.rank_zero:wrapped_fn:42 - Last 64 tokens: [SEP] but metro's definition of " urbanite " is not accurate, says matthew gandy, director of the urban laboratory at university college london. [SEP] goodrich petroleum is an independent oil and gas exploration and production company listed on the new york stock exchange. [SEP] the support from spears is to end on november 15 [SEP]
+2026-08-20 19:12:06.820 | INFO     | lightning_utilities.core.rank_zero:wrapped_fn:42 - ids: tensor([  102,  2021,  6005,  1005,  1055,  6210,  1997,  1000,  3923,  4221,
+         1000,  2003,  2025,  8321,  1010,  2758,  5487, 25957,  5149,  1010,
+         2472,  1997,  1996,  3923,  5911,  2012,  2118,  2267,  2414,  1012,
+          102,  2204, 13149, 11540,  2003,  2019,  2981,  3514,  1998,  3806,
+         8993,  1998,  2537,  2194,  3205,  2006,  1996,  2047,  2259,  4518,
+         3863,  1012,   102,  1996,  2490,  2013, 13957,  2003,  2000,  2203,
+         2006,  2281,  2321,   102])
+2026-08-20 19:12:06.820 | INFO     | lightning_utilities.core.rank_zero:wrapped_fn:42 - ==================================================
+2026-08-20 19:12:06.820 | INFO     | lightning_utilities.core.rank_zero:wrapped_fn:42 - Printing valid batch.
+2026-08-20 19:12:06.821 | INFO     | lightning_utilities.core.rank_zero:wrapped_fn:42 - Batch input_ids.shape: torch.Size([1, 128])
+2026-08-20 19:12:06.821 | INFO     | lightning_utilities.core.rank_zero:wrapped_fn:42 - First 64 tokens: [CLS] in medieval times, the townspeople from norcia were so practised at butchering pigs that they gained an in - depth knowledge of anatomy and were allowed to practise as surgeons along with members of the clergy. [SEP] the satellite weighs some 660 pounds, the israeli haaretz newspaper reported, citing unnamed company officials.
+2026-08-20 19:12:06.821 | INFO     | lightning_utilities.core.rank_zero:wrapped_fn:42 - ids: tensor([  101,  1999,  5781,  2335,  1010,  1996, 27938,  2013,  4496,  7405,
+         2020,  2061, 20439,  2012, 14998,  2075, 14695,  2008,  2027,  4227,
+         2019,  1999,  1011,  5995,  3716,  1997, 13336,  1998,  2020,  3039,
+         2000, 10975, 18908,  5562,  2004, 16804,  2247,  2007,  2372,  1997,
+         1996, 11646,  1012,   102,  1996,  5871, 21094,  2070, 20982,  7038,
+         1010,  1996,  5611,  5292, 12069,  5753,  3780,  2988,  1010,  8951,
+        13294,  2194,  4584,  1012])
+2026-08-20 19:12:06.821 | INFO     | lightning_utilities.core.rank_zero:wrapped_fn:42 - Last 64 tokens: [SEP] the situation has been further complicated by growing indications that the us treasury is preparing to make funds available for recapitalsing us banks. [SEP] as north korea's close neighbor, traditional ally and main provider of economic aid, china is widely believed to hold the key to solving the north korea conundrum. [SEP]
+2026-08-20 19:12:06.821 | INFO     | lightning_utilities.core.rank_zero:wrapped_fn:42 - ids: tensor([  102,  1996,  3663,  2038,  2042,  2582,  8552,  2011,  3652, 24936,
+         2008,  1996,  2149,  9837,  2003,  8225,  2000,  2191,  5029,  2800,
+         2005, 28667,  9331, 18400,  7741,  2149,  5085,  1012,   102,  2004,
+         2167,  4420,  1005,  1055,  2485, 11429,  1010,  3151,  9698,  1998,
+         2364, 10802,  1997,  3171,  4681,  1010,  2859,  2003,  4235,  3373,
+         2000,  2907,  1996,  3145,  2000, 13729,  1996,  2167,  4420,  9530,
+         8630,  6824,  1012,   102])
+2026-08-20 19:12:06.821 | INFO     | lightning_utilities.core.rank_zero:wrapped_fn:42 - ==================================================
+2026-08-20 19:12:06.821 | INFO     | lightning_utilities.core.rank_zero:wrapped_fn:42 - Train set length: 7102854
+2026-08-20 19:12:06.821 | INFO     | lightning_utilities.core.rank_zero:wrapped_fn:42 - Valid set length: 71979
+2026-08-20 19:12:06.839 | WARNING  | lightning_utilities.core.rank_zero:wrapped_fn:42 - Using Legacy dataloader!!!
+2026-08-20 19:12:06.839 | INFO     | lightning_utilities.core.rank_zero:wrapped_fn:42 - Loading data from: /home/ruizzhao/diffusion_baseline/data/training_cache/di4c/lm1b/lm1b_train_bs128_wrapped.dat
+2026-08-20 19:12:06.919 | WARNING  | lightning_utilities.core.rank_zero:wrapped_fn:42 - Using Legacy dataloader!!!
+2026-08-20 19:12:06.920 | WARNING  | lightning_utilities.core.rank_zero:wrapped_fn:42 - Using LM1B test set as validation set
+2026-08-20 19:12:06.920 | INFO     | lightning_utilities.core.rank_zero:wrapped_fn:42 - Loading data from: /home/ruizzhao/diffusion_baseline/data/training_cache/di4c/lm1b/lm1b_test_bs128_wrapped.dat