MLA | KDA | TPA | GDA | ResFormer | Mamba3 | DragonMimo (WIP) | tokenshift | SeeDNorm | shrink DA/GDN | gate shared across all block types |

configuration_dragon.py

@@ -3,6 +3,7 @@
 # TODO : TP (cf qwen)
 # TODO : init
 
+from typing import Optional
 import re
 
 from transformers.configuration_utils import PretrainedConfig
@@ -89,29 +90,40 @@ class DragonConfig(PretrainedConfig):
     model_type = "dragon"
     keys_to_ignore_at_inference = ["past_key_values"]
 
-    """
-    config.num_attention_heads_indexer
-        self.indexer_head_dim = config.head_dim_indexer
-        self.q_lora_rank = config.dsa_q_lora_rank
-        self.topk = config.dsa_topk
-        """
-
     def __init__(
         self,
+        mla_kv_rank: int = 128,
+        shrink_qk_da: int = 2,
+        shrink_qk_gdn: int = 2,
+        mixer_gn: bool = True,
+        kda_allow_neg_eigval: bool = False,
+        kda_num_v_heads: Optional[int] = None,
+        seednorm_wd: bool = True,
+        normalization_type: str = "rmsnorm",
+        tpa_rank: int = 2,
+        num_signal_heads_diff: Optional[int] = None,
+        scalar_proj_as_hidden_matrix: bool = True,
+        token_shift_attn: bool = False,
+        token_shift_gdn: bool = False,
+        token_conv1d_attn: bool = False,
+        token_conv1d_gdn: bool = True,
         patch_level_training: bool = False,
         patch_level_training_size: int = 4,
-        nsa_head_dim: int = 128,
         nsa_topk: int = 16,
         nsa_block_size: int = 64,
         nsa_window_size: int = 512,
-        cca_head_dim: int = 128,
         cca_seq_kernel_size: int = 4,
         rope_gdn: str = None,
         zero_centered_gate: bool = False,
         zero_centered_gate_type: int = 1,
         scalable_softmax: bool = True,
+        resformer: bool = False,
+        mamba_mimo_dim : int = 4,
+        gate_type: str = "elementwise",
+        gate_act: str = "silu",
         gate_attn: bool = False,
         gate_gdn: bool = True,
+        head_dim_gdn: Optional[int] = None,
         num_attention_heads_gdn: int = 32,
         num_key_value_heads_gdn: int = None,
         fused_loss_computation=False,
@@ -129,6 +141,7 @@ class DragonConfig(PretrainedConfig):
         intermediate_size=8192,
         expand_factor=2,
         layers_config=4*"lrdlr",
+        head_dim=128,
         num_attention_heads=32,
         num_key_value_heads=8,
         mlp_hidden_act="relu2",
@@ -147,7 +160,10 @@ class DragonConfig(PretrainedConfig):
         eos_token_id=2,
         sliding_window_size=1024,
         slw_wsize=-1,
+        rope_type_local="rope",
+        rope_type_global="",
         rope_theta_local=163.,
+        rope_theta_global=10000.,
         uscaling_tau=0.2,
         attention_dropout=0.,
         hidden_dropout=0.,
@@ -157,21 +173,39 @@ class DragonConfig(PretrainedConfig):
         gdn_dt_init_floor=1e-4,
         gdn_A_init_range=(1, 16),
         old_lns=False,
+        mlp_linking=False,
         **kwargs,
     ):
+        self.mla_kv_rank = mla_kv_rank
+        self.shrink_qk_da = shrink_qk_da
+        self.shrink_qk_gdn = shrink_qk_gdn
+        self.mixer_gn = mixer_gn
+        self.kda_allow_neg_eigval = kda_allow_neg_eigval
+        self.kda_num_v_heads = kda_num_v_heads
+        self.seednorm_wd = seednorm_wd
+        self.normalization_type = normalization_type
+        self.tpa_rank = tpa_rank
+        self.num_signal_heads_diff = num_signal_heads_diff
+        self.scalar_proj_as_hidden_matrix = scalar_proj_as_hidden_matrix
+        self.token_shift_attn = token_shift_attn
+        self.token_shift_gdn = token_shift_gdn
+        self.token_conv1d_attn = token_conv1d_attn
+        self.token_conv1d_gdn = token_conv1d_gdn
         self.patch_level_training = patch_level_training
         self.patch_level_training_size = patch_level_training_size
-        self.nsa_head_dim = nsa_head_dim
         self.nsa_topk = nsa_topk
         self.nsa_block_size = nsa_block_size
         self.nsa_window_size = nsa_window_size
-        self.cca_head_dim = cca_head_dim
         self.cca_seq_kernel_size = cca_seq_kernel_size
         self.rope_gdn = rope_gdn
         self.zero_centered_gate = zero_centered_gate
         self.zero_centered_gate_type = zero_centered_gate_type
+        self.gate_type = gate_type
+        self.gate_act = gate_act
         self.gate_attn = gate_attn
         self.gate_gdn = gate_gdn
+        self.head_dim = head_dim
+        self.head_dim_gdn = head_dim_gdn
         self.num_attention_heads_gdn = num_attention_heads_gdn
         if num_key_value_heads_gdn is None:
             num_key_value_heads_gdn = num_attention_heads_gdn
@@ -182,13 +216,18 @@ class DragonConfig(PretrainedConfig):
         self.dsa_q_lora_rank = dsa_q_lora_rank
         self.dsa_topk = dsa_topk
         self.zero_centered_gamma = zero_centered_gamma
-        self.rope_theta = rope_theta_local
+        self.rope_type_local = rope_type_local
+        self.rope_type_global = rope_type_global
+        self.rope_theta_local = rope_theta_local
+        self.rope_theta_global = rope_theta_global
         self.qk_norm = qk_norm
         self.softcap_local_attn=softcap_local_attn
         self.softcap_global_attn=softcap_global_attn
         self.use_uscaling = use_uscaling
         self.uscaling_tau = uscaling_tau
         self.scalable_softmax = scalable_softmax
+        self.resformer = resformer
+        self.mamba_mimo_dim = mamba_mimo_dim
 
         self.vocab_size = vocab_size
         self.tie_word_embeddings = tie_word_embeddings
@@ -226,9 +265,11 @@ class DragonConfig(PretrainedConfig):
         self.A_init_range = gdn_A_init_range
 
         self.old_lns = old_lns
+        
+        self.mlp_linking = mlp_linking
 
-        assert self.hidden_size % self.num_attention_heads == 0
-        assert self.num_attention_heads % self.num_key_value_heads == 0
+        #assert self.hidden_size % self.num_attention_heads == 0
+        #assert self.num_attention_heads % self.num_key_value_heads == 0
         #assert self.num_attention_heads % 2 == 0, "Number of attention heads must be even for differential attention."
         #assert self.num_key_value_heads % 2 == 0, "Number of kv heads must be even for differential attention."
 

inspecting_dragon.py

@@ -0,0 +1,302 @@
+from typing import List, Dict, Optional
+from dataclasses import dataclass
+import json
+import re
+import torch
+import torch.nn as nn
+from functools import partial
+from collections import defaultdict
+import tyro
+
+from .configuration_dragon import DragonConfig
+from .modeling_dragon import DragonForCausalLM
+
+@dataclass
+class NanoArgs:
+    resume_from: Optional[str] = None
+    run_name : str = ""
+    
+    # arch - general
+    d_model : int = 768
+    n_heads : int = 6 # head dim 128 suggested by @Grad62304977
+    layers_config : str = 4*"lrdlr"
+    expand_factor : int = 1 # expand factor for Mamba/Dragon
+    rope_theta_local: float = 10000.0
+    eps_rmsnorm: float = 1e-6
+    mlp_expand: int = 4 # expand factor for MLP
+    fused_loss_computation : bool = True # whether to use fused linear + cross entropy loss
+    use_uscaling: bool = False
+    uscaling_tau: float = 0.2
+    zero_centered_gamma: bool = False
+    zero_centered_gate: bool = False
+    zero_centered_gate_type: int = 1 # 1, 2, 3, 4
+    gate_attn: bool = False
+    gate_gdn: bool = True
+    gate_type: str = "elementwise" # elementwise (one per dim), headwise (one per head)
+    gate_act: str = "silu" # silu, sigmoid
+    scalar_proj_as_hidden_matrix: bool = True
+
+    # attention related
+    n_kv_heads : int = 0
+    swa_window_size : int = 1024
+    slw_warmup_iters: float = 0
+    slw_start: int = 8 # window size at the start of training
+    slw_increment: int = 64 # window size increment at each step
+    softcap_local_attn: float = 0.0 # logit soft-capping for local attn logits, as per Gemma2 (0.0 = no soft-capping)
+    softcap_global_attn: float = 0.0
+    qk_norm: bool = True
+    scalable_softmax: bool = True
+    token_shift: bool = False
+    num_attention_heads_indexer: int = 8
+    head_dim_indexer: int = 32
+    dsa_q_lora_rank: int = 128
+    dsa_topk: int = 512
+    cca_head_dim: int = 128
+    cca_seq_kernel_size: int = 4
+    nsa_head_dim: int = 128
+    nsa_topk: int = 16
+    nsa_block_size: int = 64
+    nsa_window_size: int = 512
+
+    # GDN related
+    rope_gdn: Optional[str] = None # None, rope, (srope)
+    n_heads_gdn: int = 0
+    n_kv_heads_gdn: int = 0
+
+    # optim
+    optim: str = "adamw" # adamw, spam, stable-spam, muon, muon_moonlight, splus
+    second_order_optim : Optional[str] = None #Snoo
+    batch_size: int = 8*64 # batch size, in sequences, across all devices
+    device_batch_size: int = 64 # batch size, in sequences, per device
+    total_iterations: int = 1000 # number of iterations to run
+    learning_rate: float = 1e-4
+    weight_decay: float = 0.
+    adam_beta1: float = 0.9
+    adam_beta2: float = 0.95
+    adam_eps: float = 1e-8
+    warmup_iters: int = 200
+    warmdown_iters: int = 3000
+    grad_norm_clip: float = 1.0
+    uscaling_mult_embed: float = 0
+    uscaling_mult_scalar: float = 0
+    uscaling_mult_head: float = 0
+    init_std: float = 0.006
+    patch_level_training: bool = False
+    patch_level_training_size: int = 4
+    patch_level_training_mode: str = "reduced" # reduced = ask L tokens, treat L//K. full = ask K*L tokens, treat L.
+
+    # data
+    vocab_size: int = 50304
+    sequence_length: int = 1024
+    use_patch_level_training: bool = False
+    patch_size: int = 4
+    patch_training_fraction: float = 0.67
+    input_bin: Optional[str] = None
+    input_val_bin: Optional[str] = None
+
+    # evaluation and logging
+    val_loss_every: int = 125
+    val_iterations: int = 50 # 1 step = global bs * T tokens
+    inspect_every: int = 0
+    save_every: int = 1000
+    log_dir: str = "logs/"
+    wandb_project: str = "dragon_v1.5"
+    wandb_name: Optional[str] = None
+    log_wandb: bool = False
+
+    load_arg_from_config: bool = True
+    load_optim: bool = True
+    load_sched: bool = True
+    compile: bool = True
+
+    # used during training
+    slw_window: int = 0
+
+args = tyro.cli(NanoArgs)
+
+# load model.
+config_hf = DragonConfig(
+    scalar_proj_as_hidden_matrix=args.scalar_proj_as_hidden_matrix,
+    token_shift=args.token_shift,
+    patch_level_training=args.patch_level_training,
+    patch_level_training_size=args.patch_level_training_size,
+    nsa_head_dim=args.nsa_head_dim,
+    nsa_topk=args.nsa_topk,
+    nsa_block_size=args.nsa_block_size,
+    nsa_window_size=args.nsa_window_size,
+    cca_head_dim=args.cca_head_dim,
+    cca_seq_kernel_size=args.cca_seq_kernel_size,
+    num_attention_heads_gdn=args.n_heads_gdn,
+    num_key_value_heads_gdn=args.n_kv_heads_gdn,
+    zero_centered_gate=args.zero_centered_gate,
+    zero_centered_gate_type=args.zero_centered_gate_type,
+    scalable_softmax=args.scalable_softmax,
+    gate_type=args.gate_type,
+    gate_act=args.gate_act,
+    gate_attn=args.gate_attn,
+    gate_gdn=args.gate_gdn,
+    fused_loss_computation=args.fused_loss_computation,
+    qk_norm=args.qk_norm,
+    num_attention_heads_indexer=args.num_attention_heads_indexer,
+    head_dim_indexer=args.head_dim_indexer,
+    dsa_q_lora_rank=args.dsa_q_lora_rank,
+    dsa_topk=args.dsa_topk,
+    zero_centered_gamma=args.zero_centered_gamma,
+    vocab_size=args.vocab_size,
+    max_position_embeddings=args.sequence_length,
+    use_uscaling=args.use_uscaling,
+    hidden_size=args.d_model,
+    intermediate_size=args.d_model * args.mlp_expand,
+    expand_factor=args.expand_factor,
+    layers_config=args.layers_config,
+    num_attention_heads=args.n_heads,
+    num_key_value_heads=args.n_kv_heads if args.n_kv_heads > 0 else args.n_heads,
+    initializer_range=args.init_std,
+    softcap_local_attn=args.softcap_local_attn,
+    softcap_global_attn=args.softcap_global_attn,
+    norm_epsilon=args.eps_rmsnorm,
+    use_cache=False,
+    sliding_window_size=args.swa_window_size,
+    rope_theta_local=args.rope_theta_local,
+    uscaling_tau=args.uscaling_tau,
+)
+
+model = DragonForCausalLM(config_hf)
+model = model.cuda()
+
+B, L = 2, 2048
+
+# ---------- helpers ---------- #
+def l1(x: torch.Tensor) -> float:
+    return x.abs().mean().item()
+
+def _capture(name: str, store: Dict[str, torch.Tensor], _m, _inp, out):
+    """Save every tensor produced by a module so that we can measure activations."""
+    def walk(x, suf=""):
+        if torch.is_tensor(x):
+            store[f"{name}{suf}"] = x.detach()
+        elif isinstance(x, (list, tuple)):
+            for i, xi in enumerate(x):
+                walk(xi, suf + f"[{i}]")
+    walk(out)
+
+_stat_pat = re.compile(r"(\.grad\.(?:std|mean|l1)|\.act\.(?:std|mean|l1)|\.(?:std|mean|l1))$")
+
+# Support multiple model naming schemes
+_LAYER_PATTERNS = [
+    re.compile(r"\.h\.(\d+)\."),                 # transformer.h.<i>.
+    re.compile(r"\.layers\.(\d+)\."),            # model.layers.<i>.
+    re.compile(r"\.decoder\.layers\.(\d+)\."),   # decoder.layers.<i>.
+    re.compile(r"\.block\.(\d+)\."),             # ...block.<i>.
+]
+
+def _find_layer_span_and_idx(key: str):
+    for pat in _LAYer_PATTERNS if False else _LAYER_PATTERNS:  # keep exact name
+        m = pat.search(key)
+        if m:
+            return m.span(0), int(m.group(1))  # span of ".layers.<i>." and the idx
+    return None, -1
+
+def _layer_idx(key: str) -> int:
+    _, idx = _find_layer_span_and_idx(key)
+    return idx
+
+def _base_key(key: str) -> str:
+    """Return <parameter-suffix>.<stat> without the layer index, e.g. mixer.linear_qkv.weight.std"""
+    span, _ = _find_layer_span_and_idx(key)
+    pre_cut = key
+    if span:
+        s, e = span
+        pre_cut = pre_cut[:s] + "." + pre_cut[e:]  # collapse the layer segment to a single dot
+    # Drop common top-level prefixes
+    for prefix in ("transformer.", "model.", "module."):
+        if pre_cut.startswith(prefix):
+            pre_cut = pre_cut[len(prefix):]
+    stat_match = _stat_pat.search(pre_cut)
+    assert stat_match, f"No stat suffix in key {key}"
+    stat_suffix = stat_match.group(1)
+    base_no_stat = pre_cut[: -len(stat_suffix)]
+    return f"{base_no_stat}{stat_suffix}"
+
+# ---------- main routine ---------- #
+
+def show_layer_stats(model: nn.Module) -> str:
+    """Run a forward/backward pass and return aggregated stats in JSON.
+
+    The JSON schema is:
+    {
+        "attn.linear_qkv.weight.std": [layer0, layer1, ..., layerN],
+        "attn.linear_qkv.grad.std"  : [...],
+        "attn.linear_qkv.act.std"   : [...],
+        ...
+    }
+    Layers that do not have a value for a given statistic are represented with null.
+    Non‑layer parameters (e.g., embeddings) are kept flat as a single key‑value pair.
+    """
+
+    PAD = len(str(len(config_hf.layers_config) - 1))
+
+    # ----- collect activations ----- #
+    acts, hooks = {}, []
+    for n, m in model.named_modules():
+        if m is model:
+            continue  # skip root
+        hooks.append(m.register_forward_hook(partial(_capture, n, acts)))
+
+    x = torch.randint(0, config_hf.vocab_size, (B, L), device="cuda")
+    y = torch.randint(0, config_hf.vocab_size, (B, L), device="cuda")
+    loss = model(input_ids=x, labels=y).loss
+    loss.backward()
+
+    # ----- collect stats (weight / grad / act) ----- #
+    raw_stats = {}
+    for n, p in model.named_parameters():
+        raw_stats[f"{n}.std"]      = p.std().item()
+        #raw_stats[f"{n}.mean"]     = p.mean().item()
+        raw_stats[f"{n}.l1"]       = l1(p)
+        if p.grad is not None:
+            raw_stats[f"{n}.grad.std"]  = p.grad.std().item()
+            #raw_stats[f"{n}.grad.mean"] = p.grad.mean().item()
+            raw_stats[f"{n}.grad.l1"]   = l1(p.grad)
+    for n, a in acts.items():
+        raw_stats[f"{n}.act.std"]  = a.std().item()
+        #raw_stats[f"{n}.act.mean"] = a.mean().item()
+        raw_stats[f"{n}.act.l1"]   = l1(a)
+
+    # ----- aggregate across layers ----- #
+    agg: Dict[str, List] = defaultdict(lambda: [None] * len(config_hf.layers_config))
+    flat: Dict[str, float] = {}
+
+    for key, val in raw_stats.items():
+        layer = _layer_idx(key)
+        if layer == -1:
+            # params without layer index stay flat
+            flat[key] = val
+            continue
+        base = _base_key(key)
+        if layer < len(config_hf.layers_config):
+            agg[base][layer] = val
+        else:
+            # unexpected layer index; fall back to flat
+            flat[key] = val
+
+    # ----- merge flat & aggregated with custom sorting ----- #
+    stats = {}
+
+    # First: per-quantity arrays over layers
+    for base_key in sorted(agg.keys()):
+        stats[f"inspect/{base_key}"] = agg[base_key]  # list of length = #layers (None where absent)
+
+    # Then: non-layer (“flat”) stats
+    for k, v in sorted(flat.items()):
+        stats[f"inspect/{k}"] = v
+    
+    return stats
+
+filename = "layer_stats.json"
+
+json_blob = show_layer_stats(model)
+with open(args.log_dir + filename, "w") as f:
+    if json_blob:
+        json.dump(json_blob, f, indent=2)  # Use json.dump() instead of f.write()
+print(f"✅ Saved layer stats to {args.log_dir + filename} ✅")

optimizers/Ademamix.py

@@ -0,0 +1,165 @@
+"""
+Adapted from: https://pytorch.org/docs/1.6.0/_modules/torch/optim/adam.html
+"""
+import math
+import torch
+from torch.optim import Optimizer
+
+
+def linear_warmup_scheduler(step, alpha_end, alpha_start=0, warmup=1):
+    if step < warmup:
+        a = step / float(warmup)
+        return (1.0-a) * alpha_start + a * alpha_end
+    return alpha_end
+
+
+def linear_hl_warmup_scheduler(step, beta_end, beta_start=0, warmup=1):
+
+    def f(beta, eps=1e-8):
+        return math.log(0.5)/math.log(beta+eps)-1
+
+    def f_inv(t):
+        return math.pow(0.5, 1/(t+1))
+
+    if step < warmup:
+        a = step / float(warmup)
+        return f_inv((1.0-a) * f(beta_start) + a * f(beta_end))
+    return beta_end
+
+
+class AdEMAMix(Optimizer):
+    r"""Implements the AdEMAMix algorithm.
+
+    Arguments:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups
+        lr (float, optional): learning rate (default: 1e-3)
+        betas (Tuple[float, float, float], optional): coefficients used for computing
+            running averages of gradient and its square (default: (0.9, 0.999, 0.9999)) 
+            corresponding to beta_1, beta_2, beta_3 in AdEMAMix
+        alpha (float): AdEMAMix alpha coeficient mixing the slow and fast EMAs (default: 2)
+        beta3_warmup (int, optional): number of warmup steps used to increase beta3 (default: None)
+        alpha_warmup: (int, optional): number of warmup steps used to increase alpha (default: None)
+        eps (float, optional): term added to the denominator to improve
+            numerical stability (default: 1e-8)
+        weight_decay (float, optional): weight decay as in AdamW (default: 0)
+    """
+
+    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999, 0.999), alpha=8.0, 
+                 beta3_warmup=None, alpha_warmup=None,  eps=1e-8,
+                 weight_decay=0):
+        if not 0.0 <= lr:
+            raise ValueError("Invalid learning rate: {}".format(lr))
+        if not 0.0 <= eps:
+            raise ValueError("Invalid epsilon value: {}".format(eps))
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
+        if not 0.0 <= betas[2] < 1.0:
+            raise ValueError("Invalid beta parameter at index 2: {}".format(betas[2]))
+        if not 0.0 <= weight_decay:
+            raise ValueError("Invalid weight_decay value: {}".format(weight_decay))
+        if not 0.0 <= alpha:
+            raise ValueError("Invalid alpha value: {}".format(alpha))
+        defaults = dict(lr=lr, betas=betas, eps=eps, alpha=alpha, beta3_warmup=beta3_warmup,
+                        alpha_warmup=alpha_warmup, weight_decay=weight_decay)
+        super(AdEMAMix, self).__init__(params, defaults)
+
+    def __setstate__(self, state):
+        super(AdEMAMix, self).__setstate__(state)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        """Performs a single optimization step.
+
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            
+            lr = group["lr"]
+            lmbda = group["weight_decay"]
+            eps = group["eps"]
+            beta1, beta2, beta3_final = group["betas"]
+            beta3_warmup = group["beta3_warmup"]
+            alpha_final = group["alpha"]
+            alpha_warmup = group["alpha_warmup"]
+        
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                grad = p.grad
+                if grad.is_sparse:
+                    raise RuntimeError('AdEMAMix does not support sparse gradients.')
+
+                state = self.state[p]
+
+                # State initialization
+                if len(state) == 0:
+                    state['step'] = 0
+                    # Exponential moving average of gradient values
+                    if beta1 != 0.0: # save memory in case beta1 is 0.0
+                        state['exp_avg_fast'] = torch.zeros_like(p, memory_format=torch.preserve_format)
+                    else: 
+                        state['exp_avg_fast'] = None
+                    state['exp_avg_slow'] = torch.zeros_like(p, memory_format=torch.preserve_format)
+                    # Exponential moving average of squared gradient values
+                    state['exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)
+
+                exp_avg_fast, exp_avg_slow, exp_avg_sq = state['exp_avg_fast'], state['exp_avg_slow'], state['exp_avg_sq']
+
+                state['step'] += 1
+                bias_correction1 = 1 - beta1 ** state['step']
+                bias_correction2 = 1 - beta2 ** state['step']
+
+                # Compute the effective alpha and beta3 in case warmup is used 
+                if alpha_warmup is not None:
+                    alpha = linear_warmup_scheduler(state["step"], alpha_end=alpha_final, alpha_start=0, warmup=alpha_warmup)
+                else:
+                    alpha = alpha_final
+                
+                if beta3_warmup is not None:
+                    beta3 = linear_hl_warmup_scheduler(state["step"], beta_end=beta3_final, beta_start=beta1, warmup=beta3_warmup)
+                else:
+                    beta3 = beta3_final
+
+                # Decay the first and second moment running average coefficient
+                if beta1 != 0.0:
+                    exp_avg_fast.mul_(beta1).add_(grad, alpha=1 - beta1)
+                else:
+                    exp_avg_fast = grad
+                exp_avg_slow.mul_(beta3).add_(grad, alpha=1 - beta3)
+                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
+
+                denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps)
+
+                update = (exp_avg_fast.div(bias_correction1) + alpha * exp_avg_slow) / denom
+
+                # decay
+                update.add_(p, alpha=lmbda)
+
+                p.add_(-lr * update)
+
+        return loss
+
+
+if __name__ == "__main__": # small dummy test
+
+    x = torch.randn((10,7))
+    model = torch.nn.Linear(7, 1, bias=False)
+    opt = AdEMAMix(params=model.parameters(), lr=1e-2, betas=(0.9, 0.999, 0.9999), alpha=2.0, beta3_warmup=45, alpha_warmup=45, weight_decay=0.1)
+    print(model.weight)
+    for itr in range(50):
+        y = model(x).mean()
+        opt.zero_grad()
+        y.backward()
+        opt.step()
+
+    print(model.weight)

optimizers/Snoo.py

@@ -0,0 +1,67 @@
+import torch
+import torch.nn as nn
+
+
+
+class Snoo:
+    """
+    @DominikKallusky, @vishal9-team, @vinaysrao
+    Sparse Nesterov Outer Optimizer (Snoo) is a momentum-based wrapper to any optimizer that can
+    improve the stability and smoothness of the optimization process and thus the quality
+    of large language models (LLM) and other models. Snoo implicitly adds temporal regularization
+    to the parameters, thus smoothing the training trajectory and instilling a bias towards flatter
+    minima and lower parameter norms. Snoo is computationally efficient, incurring minimal overhead
+    in compute and moderate memory usage.
+    """
+
+    @torch.no_grad()
+    def __init__(self, model: nn.Module, lr: float, momentum: float, k: int) -> None:
+        self.model = model
+        self.lr = lr
+        self.momentum = momentum
+        self.k = k
+        self.current_step = 0
+        self.outer_buf = [p.clone() for p in model.parameters()]
+        self.model_params = list(self.model.parameters())
+        self.optimizer = torch.optim.SGD(
+            self.model.parameters(),
+            lr=lr,
+            momentum=momentum,
+            nesterov=True,
+            fused=True,
+        )
+
+    @torch.no_grad()
+    def step(
+        self,
+    ) -> None:
+        if self.current_step % self.k == 0:
+            for p_new, p_old in zip(self.model_params, self.outer_buf):
+                p_new.grad = p_old.data - p_new.data
+                p_new.copy_(p_old, non_blocking=True)
+
+            self.optimizer.step()
+
+            for p_new, p_old in zip(self.model_params, self.outer_buf):
+                p_old.copy_(p_new, non_blocking=True)
+        self.current_step += 1
+
+    def state_dict(self):
+        state_dict = {
+            "current_step": self.current_step,
+            "lr": self.lr,
+            "momentum": self.momentum,
+            "k": self.k,
+            "outer_buf": [p.clone() for p in self.outer_buf],
+            "optimizer_state_dict": self.optimizer.state_dict(),
+        }
+        return state_dict
+
+    def load_state_dict(self, state_dict):
+        self.current_step = state_dict["current_step"]
+        self.lr = state_dict["lr"]
+        self.momentum = state_dict["momentum"]
+        self.k = state_dict["k"]
+        for p_src, p_dst in zip(state_dict["outer_buf"], self.outer_buf):
+            p_dst.copy_(p_src)
+        self.optimizer.load_state_dict(state_dict["optimizer_state_dict"])

optimizers/__init__.py

@@ -0,0 +1,2 @@
+from .Ademamix import AdEMAMix
+from .Snoo import Snoo

training_dragon.py

@@ -32,9 +32,13 @@ class NanoArgs:
     # arch - general
     d_model : int = 768
     n_heads : int = 6 # head dim 128 suggested by @Grad62304977
+    head_dim: Optional[int] = None
     layers_config : str = 4*"lrdlr"
-    expand_factor : int = 1 # expand factor for Mamba/Dragon
+    expand_factor : int = 2 # expand factor for Mamba/Dragon
+    rope_type_local: str = "rope" #p-rope
+    rope_type_global: str = "rope" #p-rope
     rope_theta_local: float = 10000.0
+    rope_theta_global: float = 0.0
     eps_rmsnorm: float = 1e-6
     mlp_expand: int = 4 # expand factor for MLP
     fused_loss_computation : bool = True # whether to use fused linear + cross entropy loss
@@ -42,9 +46,16 @@ class NanoArgs:
     uscaling_tau: float = 0.2
     zero_centered_gamma: bool = False
     zero_centered_gate: bool = False
-    zero_centered_gate_type: int = 1 # 1, 2, 3
+    zero_centered_gate_type: int = 1 # 1, 2, 3, 4
     gate_attn: bool = False
     gate_gdn: bool = True
+    gate_type: str = "elementwise" # elementwise (one per dim), headwise (one per head), kimi (lora)
+    gate_act: str = "silu" # silu, sigmoid
+    scalar_proj_as_hidden_matrix: bool = True
+    normalization_type: str = "rmsnorm" # rmsnorm, seednorm
+    seednorm_wd: bool = True
+    mixer_gn: bool = True
+    mlp_linking : bool = False
 
     # attention related
     n_kv_heads : int = 0
@@ -56,24 +67,36 @@ class NanoArgs:
     softcap_global_attn: float = 0.0
     qk_norm: bool = True
     scalable_softmax: bool = True
+    resformer : bool = False # Works only on f layers (DiffAttention)
+    token_shift_attn: bool = False
+    token_shift_gdn: bool = False
+    token_conv1d_attn: bool = False
+    token_conv1d_gdn: bool = True
     num_attention_heads_indexer: int = 8
     head_dim_indexer: int = 32
     dsa_q_lora_rank: int = 128
     dsa_topk: int = 512
-    cca_head_dim: int = 128
     cca_seq_kernel_size: int = 4
-    nsa_head_dim: int = 128
     nsa_topk: int = 16
     nsa_block_size: int = 64
     nsa_window_size: int = 512
+    num_signal_heads_diff: Optional[int] = None
+    tpa_rank: int = 2
+    shrink_qk_da: int = 2
+    mla_kv_rank: int = 128
 
     # GDN related
     rope_gdn: Optional[str] = None # None, rope, (srope)
+    head_dim_gdn: Optional[int] = None
     n_heads_gdn: int = 0
     n_kv_heads_gdn: int = 0
+    shrink_qk_gdn: int = 2
+    kda_allow_neg_eigval: bool = False
+    kda_num_v_heads: Optional[int] = None
 
     # optim
     optim: str = "adamw" # adamw, spam, stable-spam, muon, muon_moonlight, splus
+    second_order_optim : Optional[str] = None # snoo
     batch_size: int = 8*64 # batch size, in sequences, across all devices
     device_batch_size: int = 64 # batch size, in sequences, per device
     total_iterations: int = 1000 # number of iterations to run
@@ -91,14 +114,13 @@ class NanoArgs:
     init_std: float = 0.006
     patch_level_training: bool = False
     patch_level_training_size: int = 4
-    patch_level_training_mode: str = "reduced" # reduced = ask L tokens, treat L//K. full = ask K*L tokens, treat L.
+    second_order_lr: float = 0.68
+    second_order_momentum: float = 0.37
+    second_order_interval: int = 25
 
     # data
     vocab_size: int = 50304
     sequence_length: int = 1024
-    use_patch_level_training: bool = False
-    patch_size: int = 4
-    patch_training_fraction: float = 0.67
     input_bin: Optional[str] = None
     input_val_bin: Optional[str] = None
 
@@ -213,11 +235,15 @@ def param_groups_mup(model, base_lr_hidden, base_lr_scalar, base_lr_embed, base_
     for mod in model.modules():
         if isinstance(mod, nn.Linear):
             pname = id2name.get(id(mod.weight), "")
+            is_scalar = getattr(mod, "is_scalar_weight", False)
             fan_in = mod.weight.shape[1]
             scale = 1 / math.sqrt(fan_in)
             if "lm_head" in pname:
                 lr_scaled = base_lr_head
                 wd_scaled = 0.0
+            elif is_scalar:
+                lr_scaled = base_lr_scalar
+                wd_scaled = 0.0
             else:
                 lr_scaled = base_lr_hidden * scale
                 wd_scaled = wd / lr_scaled
@@ -226,7 +252,7 @@ def param_groups_mup(model, base_lr_hidden, base_lr_scalar, base_lr_embed, base_
             seen.add(mod.weight)
 
             if mod.bias is not None:
-                groups.append({"params": [mod.bias], "lr": lr_scaled, "weight_decay": 0.0})
+                groups.append({"params": [mod.bias], "lr": base_lr_scalar, "weight_decay": 0.0})
                 seen.add(mod.bias)
 
     for p in model.parameters():
@@ -235,13 +261,17 @@ def param_groups_mup(model, base_lr_hidden, base_lr_scalar, base_lr_embed, base_
         pname = id2name.get(id(p), "<unnamed>")
 
         if "embedding" in pname:
-            fan_out = p.shape[1] # nn.Embedding is transposed
+            #fan_out = p.shape[1] # nn.Embedding is transposed
             #lr_scaled = base_lr / math.sqrt(fan_out) # u-muP
             lr_scaled = base_lr_embed
         else:
             lr_scaled = base_lr_scalar
 
-        groups.append({"params": [p], "lr": lr_scaled, "weight_decay": 0.})
+        wd_scaled = 0.
+        if getattr(p, "requires_weight_decay", False):
+            wd_scaled = wd / lr_scaled
+
+        groups.append({"params": [p], "lr": lr_scaled, "weight_decay": wd_scaled})
 
     return groups
 
@@ -299,11 +329,13 @@ if master_process:
         with open(f'{logdir}/args.json', 'w') as f: json.dump(vars(args), f)
         with open(f'{logdir}/args.pkl', 'wb') as f: pickle.dump(args, f)
 def print0(s, console=True):
-    if master_process:
-        with open(logfile, "a") as f:
-            if console:
-                print(s)
-            print(s, file=f)
+    if not master_process: return
+    if console:
+        print(s)
+    try:
+        d=os.path.dirname(logfile); d and os.makedirs(d, exist_ok=True)
+        with open(logfile, "a", encoding="utf-8") as f: print(s, file=f)
+    except: pass
 if resume_dir is not None and args.load_arg_from_config:
     saved_args_path = os.path.join(os.path.dirname(resume_dir), "args.pkl")
     print0(f"Loading args from {saved_args_path}")
@@ -326,16 +358,14 @@ np.random.seed(seed)
 
 # define convenience variables.
 B, T = args.device_batch_size, args.sequence_length
+if args.patch_level_training:
+    T = args.patch_level_training_size * T
 assert args.batch_size % (B * ddp_world_size) == 0
 accumulation_steps = args.batch_size // (B * ddp_world_size)
 
 # load dataloaders.
-if args.patch_level_training:
-    if args.patch_level_training_mode == "reduced":
-        assert T % args.patch_level_training_size == 0, "sequence length must be divisible by patch level training size in reduced mode"
-        T = T
-    elif args.patch_level_training_mode == "full":
-        T = T * args.patch_level_training_size
+#if args.patch_level_training:
+#    assert T % args.patch_level_training_size == 0, "sequence length must be divisible by patch level training size in reduced mode"
 train_loader = DistributedDataLoader(args.input_bin, B, T, ddp_rank, ddp_world_size)
 val_loader = DistributedDataLoader(args.input_val_bin, B, T, ddp_rank, ddp_world_size)
 print0(f"Training DataLoader: total number of tokens: {train_loader.ntok_total} across {len(train_loader.files)} files")
@@ -343,19 +373,38 @@ print0(f"Validation DataLoader: total number of tokens: {val_loader.ntok_total}
 
 # load model.
 config_hf = DragonConfig(
+    mla_kv_rank=args.mla_kv_rank,
+    rope_gdn=args.rope_gdn,
+    shrink_qk_da=args.shrink_qk_da,
+    shrink_qk_gdn=args.shrink_qk_gdn,
+    mixer_gn=args.mixer_gn,
+    kda_allow_neg_eigval=args.kda_allow_neg_eigval,
+    kda_num_v_heads=args.kda_num_v_heads,
+    seednorm_wd=args.seednorm_wd,
+    normalization_type=args.normalization_type,
+    tpa_rank=args.tpa_rank,
+    num_signal_heads_diff=args.num_signal_heads_diff,
+    scalar_proj_as_hidden_matrix=args.scalar_proj_as_hidden_matrix,
+    token_shift_attn=args.token_shift_attn,
+    token_shift_gdn=args.token_shift_gdn,
+    token_conv1d_attn=args.token_conv1d_attn,
+    token_conv1d_gdn=args.token_conv1d_gdn,
     patch_level_training=args.patch_level_training,
     patch_level_training_size=args.patch_level_training_size,
-    nsa_head_dim=args.nsa_head_dim,
     nsa_topk=args.nsa_topk,
     nsa_block_size=args.nsa_block_size,
     nsa_window_size=args.nsa_window_size,
-    cca_head_dim=args.cca_head_dim,
     cca_seq_kernel_size=args.cca_seq_kernel_size,
+    head_dim=args.head_dim,
+    head_dim_gdn=args.head_dim_gdn,
     num_attention_heads_gdn=args.n_heads_gdn,
     num_key_value_heads_gdn=args.n_kv_heads_gdn,
     zero_centered_gate=args.zero_centered_gate,
     zero_centered_gate_type=args.zero_centered_gate_type,
     scalable_softmax=args.scalable_softmax,
+    resformer=args.resformer,
+    gate_type=args.gate_type,
+    gate_act=args.gate_act,
     gate_attn=args.gate_attn,
     gate_gdn=args.gate_gdn,
     fused_loss_computation=args.fused_loss_computation,
@@ -380,15 +429,19 @@ config_hf = DragonConfig(
     norm_epsilon=args.eps_rmsnorm,
     use_cache=False,
     sliding_window_size=args.swa_window_size,
+    rope_type_global=args.rope_type_global,
+    rope_type_local=args.rope_type_local,
+    rope_theta_global=args.rope_theta_global,
     rope_theta_local=args.rope_theta_local,
     uscaling_tau=args.uscaling_tau,
+    mlp_linking=args.mlp_linking
 )
 
 if resume_dir is None:
     model = DragonForCausalLM(config_hf)
     model = model.cuda()
 else:
-    model = DragonForCausalLM.from_pretrained(resume_dir, torch_dtype=torch.bfloat16)
+    model = DragonForCausalLM.from_pretrained(resume_dir, config=config_hf, torch_dtype=torch.bfloat16)
     model = model.cuda()
 print0(model)
 
@@ -421,12 +474,13 @@ print0(f"number of total parameters:  {num_params}")
 uncompiled_model = model
 model = torch.compile(model, dynamic=True) if args.compile else model
 model.train()
-model = DDP(model, device_ids=[ddp_local_rank])
+model = DDP(model, device_ids=[ddp_local_rank], find_unused_parameters=args.resformer)
 raw_model = model.module
 ctx = torch.amp.autocast(device_type='cuda', dtype=torch.bfloat16)
 
 # load optimizers & schedulers.
 if args.use_uscaling:
+    #assert args.optim == "adamw", "uscaling is only supported with AdamW optimizer currently"
     param_list = param_groups_mup(
         raw_model,
         base_lr_hidden=args.learning_rate,
@@ -435,9 +489,30 @@ if args.use_uscaling:
         base_lr_head=args.uscaling_mult_head*args.learning_rate if args.uscaling_mult_head > 0 else args.learning_rate,
         wd=args.weight_decay,
     )
-    optimizer = torch.optim.AdamW(param_list, betas=(args.adam_beta1, args.adam_beta2), eps=args.adam_eps)
+    if args.optim == "adamw":
+        optimizer = torch.optim.AdamW(param_list, betas=(args.adam_beta1, args.adam_beta2), eps=args.adam_eps)
+    elif args.optim == "ademamix":
+        from .optimizers.Ademamix import AdEMAMix
+        beta3_warmup = alpha_warmup = args.total_iterations
+        optimizer = AdEMAMix(param_list, beta3_warmup=beta3_warmup, alpha_warmup=alpha_warmup, weight_decay=args.weight_decay)
+    else:
+        raise ValueError(f"Unknown optimizer for unit scaling: {args.optim}")
 else:
-    optimizer = torch.optim.AdamW(raw_model.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay, betas=(args.adam_beta1, args.adam_beta2), eps=args.adam_eps)
+    if args.optim == "adamw":
+        optimizer = torch.optim.AdamW(raw_model.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay, betas=(args.adam_beta1, args.adam_beta2), eps=args.adam_eps)
+    elif args.optim == "ademamix":
+        from .optimizers.Ademamix import AdEMAMix
+
+        beta3_warmup = alpha_warmup = args.total_iterations
+        optimizer = AdEMAMix(raw_model.parameters(), lr=args.learning_rate, beta3_warmup=beta3_warmup, alpha_warmup=alpha_warmup, weight_decay=args.weight_decay)
+    else:
+        raise ValueError(f"Unknown Optimizer: {args.optim}")
+if args.second_order_optim == "snoo":
+    from .optimizers.Snoo import Snoo
+    second_order_optim = Snoo(raw_model, lr=args.second_order_lr, momentum=args.second_order_momentum, k=args.second_order_interval)
+else:
+    second_order_optim = None
+
 optimizers = [optimizer]
 
 def get_lr_wsd(num_iterations, warmup_iters, warmdown_iters, it):
@@ -478,12 +553,13 @@ WARMUP_SKIP = 10
 
 # begin training.
 train_loader.reset()
-tokenizer = transformers.AutoTokenizer.from_pretrained("openai-community/gpt2", use_fast=True) # for saving
+#tokenizer = transformers.AutoTokenizer.from_pretrained("openai-community/gpt2", use_fast=True) # for saving
+tokenizer = transformers.AutoTokenizer.from_pretrained("/leonardo_work/BOOST_LCustodi/script/training/temp/hf_models/gpt2", use_fast=True)
 x, y = train_loader.next_batch()
 
-for iter_ in range(start_iter, args.total_iterations+1):
-    last_iter = (iter_ == args.total_iterations)
-    if iter_ == WARMUP_SKIP:
+for iter_ in range(start_iter, start_iter+args.total_iterations+1):
+    last_iter = (iter_ == start_iter+args.total_iterations)
+    if iter_ == start_iter+WARMUP_SKIP:
         training_time_ms = 0
         t0 = time.perf_counter()
     to_log = {}
@@ -521,7 +597,7 @@ for iter_ in range(start_iter, args.total_iterations+1):
         model.train()
 
         # log.
-        print0(f'iteration:{iter_:0{len(str(args.total_iterations))}d}/{args.total_iterations} val_loss:{val_loss:.4f} train_time:{training_time_ms:.0f}ms')
+        print0(f'iteration:{iter_:0{len(str(start_iter+args.total_iterations))}d}/{args.total_iterations} val_loss:{val_loss:.4f} train_time:{training_time_ms:.0f}ms')
         if master_process:
             wandb.log({"val_loss": val_loss}, step=iter_)
 
@@ -530,7 +606,7 @@ for iter_ in range(start_iter, args.total_iterations+1):
         t0 = time.perf_counter()
 
     # ----------- SAVING SECTION -----------
-    if master_process and iter_ > start_iter and (last_iter or (args.save_every > 0 and iter_ % args.save_every == 0)):
+    if master_process and (last_iter or (args.save_every > 0 and iter_ % args.save_every == 0)):
         # stop the clock.
         torch.cuda.synchronize()
         training_time_ms += 1000 * (time.perf_counter() - t0)
@@ -584,14 +660,16 @@ for iter_ in range(start_iter, args.total_iterations+1):
     for opt, sched in zip(optimizers, schedulers):
         opt.step()
         sched.step()
+    if second_order_optim:
+        second_order_optim.step()
     # null those gradients.
     model.zero_grad(set_to_none=True)
 
     # ----------- LOGGING SECTION -----------
     approx_training_time_ms = training_time_ms + 1000 * (time.perf_counter() - t0)
-    avg_step_time = approx_training_time_ms / (iter_ + 1 - WARMUP_SKIP) if iter_ >= WARMUP_SKIP else 0
+    avg_step_time = approx_training_time_ms / (iter_ + 1 - WARMUP_SKIP) if iter_ >= start_iter+WARMUP_SKIP else 0
     extra = " ".join(f"{k}:{v}" for k, v in (to_log or {}).items())
-    print0(f"iteration:{iter_+1:0{len(str(args.total_iterations))}d}/{args.total_iterations} train_loss:{train_loss.item():.4f} lr: {schedulers[0].get_last_lr()[0]:.4f} train_time:{approx_training_time_ms:.0f}ms step_avg:{avg_step_time:.2f}ms {extra}")
+    print0(f"iteration:{iter_+1:0{len(str(start_iter+args.total_iterations))}d}/{args.total_iterations} train_loss:{train_loss.item():.4f} lr: {schedulers[0].get_last_lr()[0]:.4f} train_time:{approx_training_time_ms:.0f}ms step_avg:{avg_step_time:.2f}ms {extra}")
     if master_process:
         wandb.log({'train_loss': train_loss.item(), 'step_avg_time': avg_step_time, **{f'lr_{i}': sched.get_last_lr()[0] for i, sched in enumerate(schedulers)}, 'grad_norm': grad_norm.item(), **to_log}, step=iter_)