import os
import sys

# Read the current file and the kernels file code ASAP, for logging
with open(sys.argv[0], 'r') as f: 
    code = f.read()
with open(os.path.join(os.path.dirname(sys.argv[0]), 'triton_kernels.py'), 'r') as f:
    code += f"\n\n{'-'*40}\n# triton_kernels.py\n{'-'*40}\n\n" 
    code += f.read()

import copy
import glob
import math
import threading
import time
import uuid
from dataclasses import dataclass
from itertools import accumulate
from pathlib import Path
import gc

os.environ["PYTORCH_ALLOC_CONF"] = "expandable_segments:True"
import torch
import triton

torch.empty(
    1, device=f"cuda:{os.environ['LOCAL_RANK']}", requires_grad=True
).backward()  # prevents a bug on some systems
import torch._dynamo as dynamo
import torch.distributed as dist
import torch.nn.functional as F

# torch._inductor.config.coordinate_descent_tuning = True # we have banned this flag for new records because it causes compilation to take 30min
from kernels import get_kernel
from torch import Tensor, nn

from triton_kernels import XXT, ba_plus_cAA, FusedLinearReLUSquareFunction, FusedSoftcappedCrossEntropy

dynamo.config.recompile_limit = 64

# -----------------------------------------------------------------------------
# Custom operators: FP8 matmul by @YouJiacheng
# Transposed layout by @ChrisJMcCormick allows for faster gradient accumulation.

@torch.library.custom_op("nanogpt::mm_t", mutates_args=())
def mm_t_op(x: Tensor, w: Tensor, x_s: float, w_s: float, grad_s: float) -> tuple[Tensor, Tensor, Tensor]:
    """Computes y = x @ w with F8 weights stored as (in_features, out_features)."""
    @torch.compile
    def impl(x: Tensor, w: Tensor):
        assert x.is_contiguous() and w.is_contiguous()
        assert x.shape[1] == w.shape[0]  # x: (batch, in), w: (in, out)

        x_f8 = x.div(x_s).to(torch.float8_e4m3fn)
        w_f8 = w.div(w_s).to(torch.float8_e4m3fn)

        # _scaled_mm requires column-major B. w_f8 is row-major (in, out).
        # .T.contiguous().T creates a column-major view without changing logical shape.
        w_f8_col_major = w_f8.T.contiguous().T

        out = torch._scaled_mm(
            x_f8,
            w_f8_col_major,
            out_dtype=torch.bfloat16,
            scale_a=x.new_tensor(x_s, dtype=torch.float32),
            scale_b=x.new_tensor(w_s, dtype=torch.float32),
            use_fast_accum=True,
        )
        return out, x_f8, w_f8

    return impl(x, w)

@mm_t_op.register_fake
def _(x: Tensor, w: Tensor, *_):
    assert x.ndim == w.ndim == 2
    assert x.shape[1] == w.shape[0]
    assert x.device == w.device
    assert x.is_contiguous() and w.is_contiguous()
    return x @ w, x.to(torch.float8_e4m3fn), w.to(torch.float8_e4m3fn)

@torch.library.custom_op("nanogpt::mm_t_backward", mutates_args=())
def mm_t_backward_op(g: Tensor, x_f8: Tensor, w_f8: Tensor, x_s: float, w_s: float, grad_s: float) -> tuple[Tensor, Tensor]:
    @torch.compile
    def impl(grad: Tensor, x_f8: Tensor, w_f8: Tensor):
        assert grad.is_contiguous()
        
        x_scale = grad.new_tensor(x_s, dtype=torch.float32)
        w_scale = grad.new_tensor(w_s, dtype=torch.float32)
        grad_scale = grad.new_tensor(grad_s, dtype=torch.float32)
        grad_f8 = grad.div(grad_s).to(torch.float8_e5m2)
        
        # grad_x = grad @ w.T
        grad_x = torch._scaled_mm(
            grad_f8,
            w_f8.T, 
            out_dtype=torch.bfloat16,
            scale_a=grad_scale,
            scale_b=w_scale,
            use_fast_accum=False,
        )
        
        # grad_w = x.T @ grad
        # Result is (in, out), naturally matching weight storage. No final .T needed.
        grad_w = torch._scaled_mm(
            x_f8.T.contiguous(),
            grad_f8.T.contiguous().T,
            out_dtype=torch.float32,
            scale_a=x_scale,
            scale_b=grad_scale,
            use_fast_accum=False,
        )
        
        return grad_x, grad_w

    grad_x, grad_w = impl(g, x_f8, w_f8)

    return grad_x, grad_w

@mm_t_backward_op.register_fake
def _(g: Tensor, x_f8: Tensor, w_f8: Tensor, *_):
    return x_f8.to(torch.bfloat16), w_f8.to(torch.float32)

def backward_t(ctx, grad_out: Tensor, *_):
    x_f8, w_f8 = ctx.saved_tensors
    x_s, w_s, grad_s = ctx.scales
    grad_x, grad_w = torch.ops.nanogpt.mm_t_backward(
        grad_out, x_f8, w_f8, x_s, w_s, grad_s
    )
    return grad_x, grad_w, None, None, None

def setup_context_t(ctx: torch.autograd.function.FunctionCtx, inputs, output):
    *_, x_s, w_s, grad_s = inputs
    _, x_f8, w_f8 = output
    ctx.save_for_backward(x_f8, w_f8)
    ctx.scales = x_s, w_s, grad_s
    ctx.set_materialize_grads(False)

mm_t_op.register_autograd(backward_t, setup_context=setup_context_t)

# -----------------------------------------------------------------------------
# Polar Express

# Computed for num_iters=5, safety_factor=2e-2, cushion=2
polar_express_coeffs = [
    (8.156554524902461, -22.48329292557795, 15.878769915207462),
    (4.042929935166739, -2.808917465908714, 0.5000178451051316),
    (3.8916678022926607, -2.772484153217685, 0.5060648178503393),
    (3.285753657755655, -2.3681294933425376, 0.46449024233003106),
    (2.3465413258596377, -1.7097828382687081, 0.42323551169305323)
]

@torch.compile(dynamic=False, fullgraph=True) # Must use dynamic=False or else it's much slower
def polar_express(G: torch.Tensor, split_baddbmm: bool = False):
    """
    Polar Express Sign Method: https://arxiv.org/pdf/2505.16932
    by Noah Amsel, David Persson, Christopher Musco, Robert M. Gower.
    """
    X = G.bfloat16()
    if G.size(-2) > G.size(-1):
        X = X.mT

    # Ensure spectral norm is at most 1
    X = X / (X.norm(dim=(-2, -1), keepdim=True) * (1 + 2e-2) + 1e-6)

    # Allocate buffers
    X = X.contiguous()
    A = torch.empty((*X.shape[:-1], X.size(-2)), device=X.device, dtype=X.dtype)
    B = torch.empty_like(A)
    C = torch.empty_like(X)

    # Select batched vs unbatched
    if split_baddbmm:
        BX_matmul = torch.bmm if X.ndim > 2 else torch.mm
    else:
        aX_plus_BX = torch.baddbmm if X.ndim > 2 else torch.addmm

    # Perform the iterations
    for a, b, c in polar_express_coeffs:
        XXT(X, out=A)  # A = X @ X.mT
        ba_plus_cAA(A, alpha=c, beta=b, out=B)  # B = b * A + c * A @ A

        # Referencing X twice causes pytorch to make a defensive copy,
        # resulting in a cudaMemcpyAsync in baddbmm.
        # For large matrices (i.e., the mlp weights), it's faster to split
        # the operation into two kernels to avoid this.
        if split_baddbmm:
            BX_matmul(B, X, out=C)  # C = B @ X
            C.add_(X, alpha=a)      # C = C + a*X  (in-place, X only read)
        else:
            aX_plus_BX(X, B, X, beta=a, out=C)  # C = a * X + B @ X

        X, C = C, X  # Swap references to avoid unnecessary copies

    if G.size(-2) > G.size(-1):
        X = X.mT
    return X


# -----------------------------------------------------------------------------
# Combined NorMuon + Adam Optimizer

@dataclass
class ParamConfig:
    """Per-parameter configuration for NorMuonAndAdam optimizer."""
    label: str
    optim: str  # "adam" or "normuon"
    comms: str  # "none", "replicated", or "sharded"
    adam_betas: tuple[float, float] | None
    lr_mul: float
    wd_mul: float
    lr: float
    initial_lr: float
    weight_decay: float
    # Adam-specific
    eps: float | None = None
    # NorMuon-specific
    reshape: tuple | None = None
    chunk_size: int | None = None
    momentum: float | None = None
    beta2: float | None = None
    per_matrix_lr_mul: list[float] | None = None


class NorMuonAndAdam:
    """
    Combined optimizer that handles both NorMuon (for projection matrices) and 
    Adam (for embeddings/scalars/gate weights).

    Muon - MomentUm Orthogonalized by Newton-schulz

    https://kellerjordan.github.io/posts/muon/

    Muon internally runs standard SGD-momentum, and then performs an orthogonalization post-
    processing step, in which each 2D parameter's update is replaced with the nearest orthogonal
    matrix. To efficiently orthogonalize each update, Muon uses a Newton-Schulz iteration (replaced
    here with Polar Express), which has the advantage that it can be stably run in bfloat16 on the GPU.

    Muon is applied only to the projection matrices in the attention and MLP layers, and is not recommended
    for embeddings, scalars, or individual weight vectors (e.g., bias terms or gate weights). 

    Differences from standard Muon:
    - Newton-Shulz is replaced with Polar Express for the orthogonalization step
    - NorMuon adds a low-rank variance estimator similar to Adafactor. https://arxiv.org/pdf/2510.05491
    - Cautious weight decay, a gated version of decoupled weight decay
    - Mantissa tracking for precision
    
    Adam (for embeddings/scalars/gates):
    - Standard Adam with bias correction
    - Cautious weight decay

    Configuration:
    Unlike torch.optim.Optimizer, this class uses per-parameter configs from a `param_table` dict
    and does not include parameter "groups". All parameters require a .label attribute, and a 
    corresponding entry in the param_table to specify their hyperparameters (lr_mul, wd_mul, adam_betas, etc.).

    Communication and ordering:
    Gradient communication is explicitly scheduled rather than hook-driven.
    Reductions are launched in `scatter_order`, while update math and final
    gathers are executed in `work_order`. These orders are independent and
    must each contain every parameter label exactly once.

    Two communication modes are supported per parameter:
    - 'replicated': Gradients are all-reduced and each rank computes the full update.
    - 'sharded': Gradients are reduce-scattered, each rank updates its shard,
      and results are all-gathered.

    Adam parameters may be freely sharded. NorMuon operates on full matrices; sharding is 
    supported by grouping matrices into parameter banks. NorMuon parameters must have a
    `.reshape` attribute that reshapes the bank so that the leading dimension is divisible 
    by world_size.

    # Contributors include @YouJiacheng, @KonstantinWilleke, @alexrgilbert, @adricarda,
    # @tuttyfrutyee, @vdlad, @ryanyang0, @vagrawal, @varunneal, @chrisjmccormick
    """
    def __init__(self, named_params, param_table: dict, scatter_order: list, work_order: list,
                 adam_defaults: dict, normuon_defaults: dict):
        self.world_size = dist.get_world_size() if dist.is_initialized() else 1
        
        # Store defaults for each optimizer type
        self.adam_defaults = adam_defaults
        self.normuon_defaults = normuon_defaults
        self.param_table = param_table
        self.scatter_order = scatter_order
        self.work_order = work_order
        
        # Collect params by label and build config
        self.param_cfgs: dict[nn.Parameter, ParamConfig] = {}
        self.param_states: dict[nn.Parameter, dict] = {}
        self._param_by_label: dict[str, nn.Parameter] = {}
        for name, param in named_params:
            label = getattr(param, "label", None)
            assert label is not None and label in param_table  # all params must have valid label
            assert label not in self._param_by_label  # exactly one param per label
            self._param_by_label[label] = param
            self._build_param_cfg(param, label)
        
        # Assert scatter_order and work_order match present labels exactly
        present = set(self._param_by_label.keys())
        assert set(scatter_order) == present and set(work_order) == present
        
        # Handle world_size=1: overwrite comms to "none"
        if self.world_size == 1:
            for p_cfg in self.param_cfgs.values():
                p_cfg.comms = "none"
        
        # Initialize state for all params
        self._init_state()
        
        # 0-D CPU tensors to avoid recompilation
        self._step_size_t = torch.tensor(0.0, dtype=torch.float32, device="cpu")
        self._eff_wd_t = torch.tensor(0.0, dtype=torch.float32, device="cpu")
        self._eff_lr_t = torch.tensor(0.0, dtype=torch.float32, device="cpu")
        
        # Track async operations
        self._reduce_futures: dict[nn.Parameter, tuple] = {}
        
        # Embed/lm_head tying state
        self.split_embed = False
        self._lm_head_param = self._param_by_label.get("lm_head")
        self._embed_param = self._param_by_label.get("embed")
    
    def _build_param_cfg(self, param: nn.Parameter, label: str):
        """Build config for a single parameter from param_table."""
        table_entry = self.param_table[label]
        optim = table_entry["optim"]
        comms = table_entry["comms"]
        adam_betas = table_entry.get("adam_betas")
        lr_mul = table_entry.get("lr_mul", 1.0)
        wd_mul = table_entry.get("wd_mul", 1.0)
        
        if optim == "adam":
            chunk_size = param.shape[0] // self.world_size if comms == "sharded" else None
            p_cfg = ParamConfig(
                label=label,
                optim=optim,
                comms=comms,
                adam_betas=tuple(adam_betas) if adam_betas else None,
                lr_mul=lr_mul,
                wd_mul=wd_mul,
                lr=self.adam_defaults["lr"],
                initial_lr=self.adam_defaults["lr"],
                weight_decay=self.adam_defaults["weight_decay"],
                eps=self.adam_defaults["eps"],
                chunk_size=chunk_size,
            )
        elif optim == "normuon":
            reshape = getattr(param, "reshape", None)
            if reshape is None:
                raise ValueError(f"NorMuon param {label} must have .reshape attribute")
            if reshape[0] % self.world_size != 0:
                raise ValueError(f"reshape[0]={reshape[0]} must be divisible by world_size")
            
            chunk_size = reshape[0] // self.world_size
            chunk_shape = (chunk_size, *reshape[1:])
            # Shape-based LR multiplier for NorMuon
            shape_mult = max(1.0, chunk_shape[-2] / chunk_shape[-1]) ** 0.5 if len(chunk_shape) >= 2 else 1.0
            lr_mul = shape_mult * lr_mul
            
            # Per-matrix LR multipliers for MLP c_proj (2x LR on odd indices)
            per_matrix_lr_mul = None
            if label == "mlp":
                rank = dist.get_rank() if dist.is_initialized() else 0
                start_idx = rank * chunk_size
                per_matrix_lr_mul = []
                for i in range(chunk_size):
                    global_idx = start_idx + i
                    is_c_proj = (global_idx % 2 == 1)
                    per_matrix_lr_mul.append(2.0 if is_c_proj else 1.0)
            
            p_cfg = ParamConfig(
                label=label,
                optim=optim,
                comms=comms,
                adam_betas=tuple(adam_betas) if adam_betas else None,
                lr_mul=lr_mul,
                wd_mul=wd_mul,
                lr=self.normuon_defaults["lr"],
                initial_lr=self.normuon_defaults["lr"],
                weight_decay=self.normuon_defaults["weight_decay"],
                reshape=reshape,
                chunk_size=chunk_size,
                momentum=self.normuon_defaults["momentum"],
                beta2=self.normuon_defaults["beta2"],
                per_matrix_lr_mul=per_matrix_lr_mul,
            )
        else:
            raise ValueError(f"Unknown optim type: {optim}")
        
        self.param_cfgs[param] = p_cfg
    
    def _init_state(self):
        """Initialize optimizer state for all parameters."""
        for param, p_cfg in self.param_cfgs.items():
            if p_cfg.optim == "adam":
                # Sharded params use chunk state, replicated use full state
                if p_cfg.comms == "sharded":
                    chunk = param[:p_cfg.chunk_size]
                else:
                    chunk = param
                exp_avg = torch.zeros_like(chunk, dtype=torch.float32, device=param.device)
                self.param_states[param] = dict(step=0, exp_avg=exp_avg, exp_avg_sq=torch.zeros_like(exp_avg))
            
            elif p_cfg.optim == "normuon":
                chunk_shape = (p_cfg.chunk_size, *p_cfg.reshape[1:])
                
                # Momentum buffer (FP32 for precision)
                momentum_buffer = torch.zeros(
                    chunk_shape, dtype=torch.float32, device=param.device
                )
                
                # Second momentum buffer - reduced along one dimension
                if chunk_shape[-2] >= chunk_shape[-1]:
                    second_mom_shape = (*chunk_shape[:-1], 1)
                else:
                    second_mom_shape = (*chunk_shape[:-2], 1, chunk_shape[-1])
                second_momentum_buffer = torch.zeros(
                    second_mom_shape, dtype=torch.float32, device=param.device
                )
                
                # Mantissa buffer for precision tracking
                mantissa = torch.zeros(
                    chunk_shape, dtype=torch.uint16, device=param.device
                )
                
                self.param_states[param] = dict(
                    momentum_buffer=momentum_buffer,
                    second_momentum_buffer=second_momentum_buffer,
                    mantissa=mantissa,
                )

    # -----------------------------------
    # Reduce/Gather operations
    
    def _launch_reduce(self, param: nn.Parameter, grad: Tensor):
        """Launch async reduce for a parameter based on its comms policy."""
        p_cfg = self.param_cfgs[param]
        
        if p_cfg.comms == "none":
            if p_cfg.optim == "normuon":
                # NorMuon needs reshaped gradient even without communication
                grad = grad.view(p_cfg.reshape)
            self._reduce_futures[param] = (None, grad)
        elif p_cfg.comms == "replicated":
            future = dist.all_reduce(grad, op=dist.ReduceOp.AVG, async_op=True).get_future()
            self._reduce_futures[param] = (future, grad)
        elif p_cfg.comms == "sharded":
            if p_cfg.optim == "normuon":
                # NorMuon: reshape before reduce_scatter
                grad_reshaped = grad.view(p_cfg.reshape)
                grad_chunk = torch.empty(
                    (p_cfg.chunk_size, *grad_reshaped.shape[1:]),
                    dtype=grad.dtype,
                    device=grad.device
                )
                future = dist.reduce_scatter_tensor(
                    grad_chunk, grad_reshaped.contiguous(), op=dist.ReduceOp.AVG, async_op=True
                ).get_future()
                self._reduce_futures[param] = (future, grad_chunk)
            else:
                # Adam: simple reduce_scatter
                grad_chunk = torch.empty_like(grad[:p_cfg.chunk_size])
                future = dist.reduce_scatter_tensor(
                    grad_chunk, grad, op=dist.ReduceOp.AVG, async_op=True
                ).get_future()
                self._reduce_futures[param] = (future, grad_chunk)

    def _launch_gather(self, param: nn.Parameter, p_slice: Tensor) -> "torch.futures.Future":
        """Launch async all_gather for a sharded parameter."""
        p_cfg = self.param_cfgs[param]
        if p_cfg.optim == "normuon":
            full_param = param.data.view(p_cfg.reshape)
            assert full_param.is_contiguous()
            return dist.all_gather_into_tensor(
                full_param, p_slice.contiguous(), async_op=True
            ).get_future()
        else:
            return dist.all_gather_into_tensor(
                param, p_slice.contiguous(), async_op=True
            ).get_future()

    # -----------------------------------
    # State management
    
    def reset(self):
        """Reset NorMuon momentum buffers and split_embed state (called on training reset)."""
        self.split_embed = False
        for param, p_cfg in self.param_cfgs.items():
            if p_cfg.optim == "normuon":
                p_state = self.param_states[param]
                p_state["momentum_buffer"].zero_()
                p_state["mantissa"].zero_()
                p_state["second_momentum_buffer"].zero_()
    
    def copy_lm_state_to_embed(self):
        """
        Copy the optimizer state from the lm_head to the embed at the untie point.
        This requires an all-gather + reshard because of different sharding:
        - lm_head (768, 50304) is sharded to (96, 50304) per rank (along model_dim)
        - embed (50304, 768) is sharded to (6288, 768) per rank (along vocab_size)
        
        We all-gather the lm_head momentum, transpose it, then each rank takes their
        embed shard to get the correct momentum state.
        """
        lm_head = self._lm_head_param
        embed = self._embed_param        
        lm_state = self.param_states[lm_head]
        embed_state = self.param_states[embed]
        lm_cfg = self.param_cfgs[lm_head]
        embed_cfg = self.param_cfgs[embed]
        
        embed_state['step'] = lm_state['step'] # Preserve step count for bias correction        
        
        # Copy optimizer state with all-gather + transpose + reshard
        if self.world_size > 1:
            rank = dist.get_rank()
            lm_chunk_size = lm_cfg.chunk_size  # 96
            embed_chunk_size = embed_cfg.chunk_size  # 6288
            
            # All-gather lm_head momentum to get full (768, 50304) tensor
            for key in ["exp_avg", "exp_avg_sq"]:
                lm_chunk = lm_state[key]  # (96, 50304)
                full_lm = torch.empty(lm_head.shape[0], lm_head.shape[1], dtype=lm_chunk.dtype, device=lm_chunk.device)
                dist.all_gather_into_tensor(full_lm, lm_chunk.contiguous())
                embed_state[key].copy_(full_lm.T[rank * embed_chunk_size:(rank + 1) * embed_chunk_size])
        else:
            # Single GPU: simple transpose
            for key in ["exp_avg", "exp_avg_sq"]:
                embed_state[key].copy_(lm_state[key].T)
        
        # Mark as split
        self.split_embed = True
    
    def state_dict(self):
        """Return the optimizer state as a dict."""
        return {
            "param_states": {id(p): s for p, s in self.param_states.items()},
            "param_cfgs": {id(p): s for p, s in self.param_cfgs.items()},
        }
    
    def load_state_dict(self, state_dict):
        """Load optimizer state from a dict."""
        # Build id->param mapping
        id_to_param = {id(p): p for p in self.param_cfgs.keys()}
        
        # Load state, preserving dtypes
        for param_id, saved_p_state in state_dict["param_states"].items():
            if param_id in id_to_param:
                param = id_to_param[param_id]
                p_state = self.param_states[param]
                for k, v in saved_p_state.items():
                    if isinstance(v, torch.Tensor) and k in p_state:
                        target_dtype = p_state[k].dtype
                        p_state[k] = v.to(dtype=target_dtype, device=p_state[k].device)
                    else:
                        p_state[k] = v

    # -----------------------------------
    # Unified optimizer step with explicit ordering

    @torch.no_grad()
    def step(self, do_adam: bool = True):
        """
        Combined optimizer step with explicit ordering.
        
        Args:
            do_adam: If True, update Adam params. NorMuon params always updated.
        
        Flow:
        1. Scatter phase: Launch reduces in scatter_order
        2. Work phase: Process updates in work_order
           - Wait for reduce, compute update, launch gather
        3. Finalize phase: Wait for gathers
        
        While the embeddings are tied:
        - Comms and update math are only done on lm_head.
        - We add embed.grad.T into lm_head.grad before comms.
        - After lm_head gather, we copy lm_head.data.T --> embed.data        
        """
        rank = dist.get_rank() if dist.is_initialized() else 0
        lm_param, embed_param = self._lm_head_param, self._embed_param
        
        # ===== Phase 1: Launch reduces in scatter_order =====
        for label in self.scatter_order:
            param = self._param_by_label[label]
            p_cfg = self.param_cfgs[param]
            
            if p_cfg.optim == "adam" and not do_adam:
                continue
            if param.grad is None:
                continue
            
            # lm_head when tied: aggregate embed.grad.T (transposed shapes)
            if label == "lm_head" and do_adam and not self.split_embed:
                if embed_param is not None and embed_param.grad is not None:
                    param.grad.add_(embed_param.grad.T)
            
            # Skip embed when tied (copied from lm_head after gather)
            if label == "embed" and not self.split_embed:
                continue
            
            self._launch_reduce(param, param.grad)
        
        # ===== Phase 2: Process updates in work_order =====
        gather_futures = []
        lm_head_gather_future = None
        
        for label in self.work_order:
            param = self._param_by_label[label]
            if param not in self._reduce_futures:
                continue
            
            p_cfg = self.param_cfgs[param]
            if p_cfg.optim == "adam" and not do_adam:
                continue
            # Wait for reduce
            future, grad_chunk = self._reduce_futures[param]
            if future is not None:
                future.wait()
            # Apply update based on optim type
            if p_cfg.optim == "adam":
                p_slice = self._adam_update(param, grad_chunk, p_cfg, rank)
            else:
                p_slice = self._normuon_update(param, grad_chunk, p_cfg, rank)
            # Launch gather for sharded params
            if p_cfg.comms == "sharded" and self.world_size > 1:
                gather_fut = self._launch_gather(param, p_slice)
                if label == "lm_head":
                    lm_head_gather_future = gather_fut
                else:
                    gather_futures.append(gather_fut)
        
        # ===== Phase 3: Wait for gathers, sync embed if tied =====
        # Wait for lm_head gather first so we can copy to embed while other gathers complete
        if lm_head_gather_future is not None:
            lm_head_gather_future.wait()
        
        # When tied: copy lm_head.T to embed
        if do_adam and not self.split_embed and embed_param is not None and lm_param is not None:
            embed_param.data.copy_(lm_param.data.T)
        
        # Wait for remaining gathers
        for fut in gather_futures:
            fut.wait()
        
        self._reduce_futures.clear()
        
        # Clear grads for updated params
        for param, p_cfg in self.param_cfgs.items():
            if p_cfg.optim == "adam" and not do_adam:
                continue  # Don't clear Adam grads on even steps
            param.grad = None

    # -----------------------------------
    # Adam update

    def _adam_update(self, param: nn.Parameter, grad_chunk: Tensor, p_cfg: ParamConfig, rank: int) -> Tensor:
        """Apply Adam update to a parameter. Returns the updated p_slice."""
        beta1, beta2 = p_cfg.adam_betas
        lr = p_cfg.lr * p_cfg.lr_mul
        
        # Get parameter slice
        if p_cfg.comms == "sharded":
            p_slice = param[rank * p_cfg.chunk_size:(rank + 1) * p_cfg.chunk_size]
        else:
            p_slice = param
        
        p_state = self.param_states[param]
        p_state["step"] += 1
        t = p_state["step"]
        
        bias1, bias2 = 1 - beta1 ** t, 1 - beta2 ** t
        self._step_size_t.fill_(lr * (bias2 ** 0.5 / bias1))
        self._eff_wd_t.fill_(lr * lr * p_cfg.weight_decay * p_cfg.wd_mul)
        
        NorMuonAndAdam._adam_update_step(
            p_slice, grad_chunk, p_state["exp_avg"], p_state["exp_avg_sq"],
            beta1, beta2, p_cfg.eps, self._step_size_t, self._eff_wd_t
        )
        
        return p_slice

    @staticmethod
    @torch.compile(dynamic=False, fullgraph=True)
    def _adam_update_step(p_slice, g_slice, exp_avg, exp_avg_sq, beta1, beta2, eps, step_size_t, eff_wd_t):
        """Compiled Adam update step."""
        exp_avg.mul_(beta1).add_(g_slice, alpha=1 - beta1)
        exp_avg_sq.mul_(beta2).addcmul_(g_slice, g_slice, value=1 - beta2)
        update = exp_avg.div(exp_avg_sq.sqrt().add_(eps)).mul_(step_size_t)
        # Cautious weight decay
        mask = (update * p_slice) > 0
        update.addcmul_(p_slice, mask, value=eff_wd_t)
        p_slice.add_(other=update, alpha=-1.0)

    # -----------------------------------
    # NorMuon update

    def _normuon_update(self, param: nn.Parameter, grad_chunk: Tensor, p_cfg: ParamConfig, rank: int) -> Tensor:
        """Apply NorMuon update to a parameter. Returns the updated p_slice."""
        chunk_shape = grad_chunk.shape
        
        p_state = self.param_states[param]
        grad_chunk = grad_chunk.float()  # FP32 for momentum
        
        # Momentum update
        momentum_buffer = p_state["momentum_buffer"]
        momentum_buffer.lerp_(grad_chunk, 1 - p_cfg.momentum)
        updated_grads = grad_chunk.lerp_(momentum_buffer, p_cfg.momentum)
        
        self._eff_lr_t.fill_(p_cfg.lr_mul * p_cfg.lr)
        self._eff_wd_t.fill_(p_cfg.wd_mul * p_cfg.weight_decay * p_cfg.lr)
        
        # Polar Express orthogonalization
        is_large_matrix = chunk_shape[-2] > 1024
        v_chunk = polar_express(updated_grads, split_baddbmm=is_large_matrix)
        
        # Variance reduction
        red_dim = -1 if chunk_shape[-2] >= chunk_shape[-1] else -2
        v_chunk = NorMuonAndAdam._apply_normuon_variance_reduction(
            v_chunk, p_state["second_momentum_buffer"], p_cfg.beta2, red_dim
        )
        
        # Update parameter, in place, with cautious weight decay
        param_view = param.data.view(p_cfg.reshape)
        p_slice = param_view[rank * p_cfg.chunk_size:(rank + 1) * p_cfg.chunk_size]
        
        # MLP has per-matrix LR multipliers (c_proj gets 2x LR)
        if p_cfg.per_matrix_lr_mul is not None:
            for mat_idx in range(p_cfg.chunk_size):
                self._eff_lr_t.fill_(p_cfg.lr_mul * p_cfg.per_matrix_lr_mul[mat_idx] * p_cfg.lr)
                self._eff_wd_t.fill_(p_cfg.wd_mul * p_cfg.weight_decay * p_cfg.lr)
                NorMuonAndAdam._cautious_wd_and_update_inplace(
                    p_slice[mat_idx].view(torch.uint16), p_state["mantissa"][mat_idx], v_chunk[mat_idx],
                    self._eff_wd_t, self._eff_lr_t
                )
        else:
            NorMuonAndAdam._cautious_wd_and_update_inplace(
                p_slice.view(torch.uint16), p_state["mantissa"], v_chunk,
                self._eff_wd_t, self._eff_lr_t
            )
        
        return p_slice

    @staticmethod
    @torch.compile(dynamic=False, fullgraph=True)
    def _cautious_wd_and_update_inplace(p, mantissa, grad, wd_tensor, lr_tensor):
        """
        Cautious weight decay + parameter update. wd_tensor and lr_tensor are 0-D CPU tensors.
        Mantissa is tracked to enable higher precision updates on bfloat16 parameters.
        bfloat16 format: 1 sign bit + 8 exponent bits + 7 mantissa bits = 16 bits total
        float32 format: 1 sign bit + 8 exponent bits + 23 mantissa bits = 32 bits total
        """
        assert p.dtype == mantissa.dtype == torch.uint16
        grad = grad.float()
        wd_factor = wd_tensor.to(torch.float32)
        lr_factor = lr_tensor.to(torch.float32)
        p_precise_raw = (p.to(torch.uint32) << 16) | mantissa.to(torch.uint32)
        p_precise = p_precise_raw.view(torch.float32)
        mask = (grad * p_precise) >= 0
        p_precise.copy_(p_precise - (p_precise * mask * wd_factor * lr_factor) - (grad * lr_factor))
        p.copy_((p_precise_raw >> 16).to(torch.uint16))
        mantissa.copy_(p_precise_raw.to(torch.uint16))

    @staticmethod
    @torch.compile(dynamic=False, fullgraph=True)
    def _apply_normuon_variance_reduction(v_chunk, second_momentum_buffer, beta2, red_dim):
        """NorMuon variance reduction. Algebraically fuses the normalization steps to minimize memory ops."""
        v_mean = v_chunk.float().square().mean(dim=red_dim, keepdim=True)
        red_dim_size = v_chunk.size(red_dim)
        v_norm_sq = v_mean.sum(dim=(-2, -1), keepdim=True).mul_(red_dim_size)
        v_norm = v_norm_sq.sqrt_()
        second_momentum_buffer.lerp_(v_mean.to(dtype=second_momentum_buffer.dtype), 1 - beta2)
        step_size = second_momentum_buffer.clamp_min(1e-10).rsqrt_()
        scaled_sq_sum = (v_mean * red_dim_size) * step_size.float().square()
        v_norm_new = scaled_sq_sum.sum(dim=(-2, -1), keepdim=True).sqrt_()
        final_scale = step_size * (v_norm / v_norm_new.clamp_min_(1e-10))
        return v_chunk.mul_(final_scale.type_as(v_chunk))

# -----------------------------------------------------------------------------
# PyTorch nn.Module definitions for the model

def norm(x: Tensor):
    return F.rms_norm(x, (x.size(-1),))


class CastedLinearT(nn.Module):
    """
    Linear layer with transposed weight storage (in_features, out_features) which
    addresses the slow kernel that was used for gradient accumulation. @chrisjmccormick
    """
    def __init__(self, in_features: int, out_features: int, use_fp8=False, x_s=1.0, w_s=1.0, grad_s=1.0):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.use_fp8 = use_fp8
        self.x_s = x_s
        self.w_s = w_s
        self.grad_s = grad_s
        
        self.weight = nn.Parameter(torch.empty(in_features, out_features, dtype=torch.bfloat16))
        self.reset_parameters()

    def reset_parameters(self) -> None:
        with torch.no_grad():
            nn.init.zeros_(self.weight) # @Grad62304977 and others

    def forward(self, x: Tensor):
        if self.use_fp8 and self.training:
            _x = x.flatten(0, -2)
            out = torch.ops.nanogpt.mm_t(_x, self.weight, x_s=self.x_s, w_s=self.w_s, grad_s=self.grad_s)[0]
            return out.reshape(*x.shape[:-1], -1)
        else:
            return x @ self.weight.type_as(x)

# -----------------------------------------------------------------------------
# PyTorch nn.Module definitions for the model

class Yarn(nn.Module):
    def __init__(self, head_dim, max_seq_len):
        super().__init__()
        self.head_dim = head_dim
        self.max_seq_len = max_seq_len
        self.reset()

    def rotary(self, x_BTHD):
        assert self.factor1.size(0) >= x_BTHD.size(-3)
        factor1, factor2 = (
            self.factor1[None, : x_BTHD.size(-3), None, :],
            self.factor2[None, : x_BTHD.size(-3), None, :],
        )
        x_flip = x_BTHD.view(*x_BTHD.shape[:-1], x_BTHD.shape[-1] // 2, 2).flip(-1).view(x_BTHD.shape)
        return factor1 * x_BTHD + factor2 * x_flip

    def reset(self):
        angular_freq = (1 / 1024) ** torch.linspace(0, 1, steps=self.head_dim//4, dtype=torch.float32, device=device)
        angular_freq = angular_freq.repeat_interleave(2)
        # half-truncate RoPE by @YouJiacheng (w/ base freq tuning)
        angular_freq = torch.cat([angular_freq, angular_freq.new_zeros(self.head_dim//2)])
        t = torch.arange(2*self.max_seq_len, dtype=torch.float32, device=device)
        theta = torch.outer(t, angular_freq)
        self.factor1 = nn.Buffer(
            theta.cos().to(torch.bfloat16), persistent=False
        )
        self.factor2 = nn.Buffer(
            theta.sin().to(torch.bfloat16), persistent=False
        )
        self.factor2[..., 1::2] *= -1
        self.angular_freq = angular_freq
        # start with 0.1, inspired by 0.12 from @leloykun and learnable scalars used by @brendanh0gan https://x.com/hi_tysam/status/1879693583898591283
        self.attn_scale = 0.1

    def apply(self, old_window: int, new_window: int, alpha: int=1, beta: int=32):
        rotations = args.block_size * old_window * self.angular_freq / (2 * torch.pi)
        scaling_factor = old_window / new_window
        interpolation_weight = torch.clamp((rotations - alpha) / (beta - alpha), 0, 1)
        self.angular_freq *= scaling_factor + interpolation_weight * (1 - scaling_factor)
        t = torch.arange(2*self.max_seq_len, dtype=torch.float32, device=self.angular_freq.device)
        theta = torch.outer(t, self.angular_freq)
        self.factor1.copy_(theta.cos())
        self.factor2.copy_(theta.sin())
        self.factor2[..., 1::2] *= -1
        self.attn_scale *= 0.2 * math.log(new_window / old_window) + 1

class YarnPairedHead(nn.Module):
    def __init__(self, head_dim, max_seq_len):
        super().__init__()
        self.head_dim = head_dim
        self.max_seq_len = max_seq_len
        self.reset()

    def rotary(self, x_BTHD):
        assert self.factor1.size(0) >= x_BTHD.size(-3)
        factor1, factor2 = (
            self.factor1[None, : x_BTHD.size(-3), None, :],
            self.factor2[None, : x_BTHD.size(-3), None, :],
        )
        x_flip = x_BTHD.view(*x_BTHD.shape[:-1], x_BTHD.shape[-1] // 2, 2).flip(-1).view(x_BTHD.shape)
        return factor1 * x_BTHD + factor2 * x_flip

    def reset(self):
        angular_freq = (1 / 1024) ** torch.linspace(0, 1, steps=self.head_dim//4, dtype=torch.float32, device=device)
        angular_freq = angular_freq.repeat_interleave(2)
        angular_freq = torch.cat([angular_freq, angular_freq.new_zeros(self.head_dim//2)])
        t = torch.arange(2*self.max_seq_len, dtype=torch.float32, device=device)
        t_even = 2 * t
        t_odd = 2 * t + 1
        theta1 = torch.outer(t_even, angular_freq)
        theta2 = torch.outer(t_odd, angular_freq)
        self.factor1 = nn.Buffer(
            torch.cat((theta1.cos(),theta2.cos()), dim=-1).to(torch.bfloat16), 
            persistent=False
        )
        self.factor2 = nn.Buffer(
            torch.cat((theta1.sin(),theta2.sin()), dim=-1).to(torch.bfloat16), 
            persistent=False
        )
        self.factor2[..., 1::2] *= -1
        self.angular_freq = angular_freq
        # start with 0.1, inspired by 0.12 from @leloykun and learnable scalars used by @brendanh0gan https://x.com/hi_tysam/status/1879693583898591283
        self.attn_scale = 0.1

    def apply(self, old_window: int, new_window: int, alpha: int=1, beta: int=32):
        rotations = args.block_size * old_window * self.angular_freq / (2 * torch.pi)
        scaling_factor = old_window / new_window
        interpolation_weight = torch.clamp((rotations - alpha) / (beta - alpha), 0, 1)
        self.angular_freq *= scaling_factor + interpolation_weight * (1 - scaling_factor)
        t = torch.arange(2*self.max_seq_len, dtype=torch.float32, device=self.angular_freq.device)
        t_even = 2 * t
        t_odd = 2 * t + 1
        theta1 = torch.outer(t_even, self.angular_freq)
        theta2 = torch.outer(t_odd, self.angular_freq)
        self.factor1.copy_(torch.cat((theta1.cos(),theta2.cos()), dim=-1))
        self.factor2.copy_( torch.cat((theta1.sin(),theta2.sin()), dim=-1))
        self.factor2[..., 1::2] *= -1
        self.attn_scale *= 0.2 * math.log(new_window / old_window) + 1

@dataclass
class AttnArgs:
    ve: torch.Tensor
    sa_lambdas: torch.Tensor
    seqlens: torch.Tensor
    bm_size: int
    yarn: Yarn
    key_offset: bool
    attn_gate_w: torch.Tensor
    ve_gate_w: torch.Tensor

flash_attn_interface = get_kernel('varunneal/flash-attention-3').flash_attn_interface

class CausalSelfAttention(nn.Module):
    def __init__(self, dim: int, head_dim: int, num_heads: int):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.dim = dim
        self.hdim = num_heads * head_dim
        assert self.hdim == self.dim, "num_heads * head_dim must equal model_dim"
        # Weights are stored in parameter banks and passed via forward()

    def forward(self, x: Tensor, attn_args: AttnArgs, qkvo_w: Tensor):
        B, T = x.size(0), x.size(1) # batch size, sequence length
        assert B == 1, "varlen sequences requires B == 1"
        assert T % 16 == 0
        # unpack attention args
        yarn = attn_args.yarn
        ve, sa_lambdas, key_offset = attn_args.ve, attn_args.sa_lambdas, attn_args.key_offset
        seqlens, bm_size = attn_args.seqlens, attn_args.bm_size
        # sparse gated attention to enable context based no-op by @classiclarryd
        # only include gates on layers with value embeds used on forward pass
        attn_gate_w, ve_gate_w = attn_args.attn_gate_w, attn_args.ve_gate_w

        q, k, v = F.linear(x, sa_lambdas[0] * qkvo_w[:self.dim * 3].type_as(x)).view(B, T, 3 * self.num_heads, self.head_dim).chunk(3, dim=-2)
        q, k = norm(q), norm(k) # QK norm @Grad62304977
        q, k = yarn.rotary(q), yarn.rotary(k)
        if key_offset:
            # shift keys forward for the stationary head dims. Enables 1-layer induction.
            k[:, 1:, :, self.head_dim // 2:] = k[:, :-1, :, self.head_dim // 2:]
        if ve is not None:
            ve_gate_out = 2 * torch.sigmoid(F.linear(x[..., :12], ve_gate_w)).view(B, T, self.num_heads, 1)
            v = v + ve_gate_out * ve.view_as(v) # @ KoszarskyB & @Grad62304977

        max_len = args.train_max_seq_len if self.training else (args.val_batch_size // (grad_accum_steps * world_size))

        # use flash_attn over flex_attn @varunneal. flash_attn_varlen suggested by @YouJiacheng
        y = flash_attn_interface.flash_attn_varlen_func(q[0], k[0], v[0], cu_seqlens_q=seqlens, cu_seqlens_k=seqlens,
                                                        max_seqlen_q=max_len, max_seqlen_k=max_len,
                                                        causal=True, softmax_scale=yarn.attn_scale, window_size=(bm_size, 0))
        y = y.view(B, T, self.num_heads, self.head_dim)
        y = y * torch.sigmoid(F.linear(x[..., :12], attn_gate_w)).view(B, T, self.num_heads, 1)
        y = y.contiguous().view(B, T, self.num_heads * self.head_dim) # re-assemble all head outputs side by side
        y = F.linear(y, sa_lambdas[1] * qkvo_w[self.dim * 3:].type_as(y))  # sa_lambdas[1] pre-multiplied to O @shenberg
        return y

class PairedHeadCausalSelfAttention(nn.Module):
    """
    Pairs up attention heads such that queries from head 1 can attend to keys in head 2, and vice-versa.
    Implemented by interleaving the k, q, and v for pairs of heads to form twice as long sequences
    EG [k1_h1, k2_h1, k3_h1], [k1_h2, k2_h2, k3_h2] -> [k1_h1, k1_h2, k2_h1, k2_h2, k3_h1, k3_h2], repeat for q and v
    """
    def __init__(self, dim: int, head_dim: int, num_heads: int):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.dim = dim
        self.hdim = num_heads * head_dim
        assert self.hdim == self.dim, "num_heads * head_dim must equal model_dim"
        # Weights are stored in parameter banks and passed via forward()

    def forward(self, x: Tensor, attn_args: AttnArgs, qkvo_w: Tensor):
        B, T = x.size(0), x.size(1) # batch size, sequence length
        assert B == 1, "varlen sequences requires B == 1"
        assert T % 16 == 0
        # unpack attention args
        yarn = attn_args.yarn
        ve, sa_lambdas = attn_args.ve, attn_args.sa_lambdas
        seqlens, bm_size = attn_args.seqlens, attn_args.bm_size
        attn_gate_w, ve_gate_w = attn_args.attn_gate_w, attn_args.ve_gate_w

        q, k, v = F.linear(x, sa_lambdas[0] * qkvo_w[:self.dim * 3].type_as(x)).view(B, T, 3 * self.num_heads, self.head_dim).chunk(3, dim=-2)
        q, k = norm(q), norm(k)

        # delay q,k reshape until rotary makes data contiguous, to enable view (non-copy)
        q = q.view(B, T, self.num_heads // 2, self.head_dim * 2)
        k = k.view(B, T, self.num_heads // 2, self.head_dim * 2)
        v = v.reshape(B, T*2, self.num_heads//2, self.head_dim)

        q, k = yarn.rotary(q), yarn.rotary(k)

        q = q.view(B, T*2, self.num_heads//2, self.head_dim)
        k = k.view(B, T*2, self.num_heads//2, self.head_dim)

        if ve is not None:
            ve_gate_out = 2 * torch.sigmoid(F.linear(x[..., :12], ve_gate_w)).view(B, T*2, self.num_heads//2, 1)
            v = v + ve_gate_out * ve.view_as(v)

        max_len = args.train_max_seq_len if self.training else (args.val_batch_size // (grad_accum_steps * world_size))

        # paired head correction
        seqlens = 2 * seqlens
        max_len = 2 * max_len

        y = flash_attn_interface.flash_attn_varlen_func(q[0], k[0], v[0], cu_seqlens_q=seqlens, cu_seqlens_k=seqlens,
                                                        max_seqlen_q=max_len, max_seqlen_k=max_len,
                                                        causal=True, softmax_scale=yarn.attn_scale, window_size=(bm_size, 0))
        y = y.view(B, T, self.num_heads, self.head_dim)
        y = y * torch.sigmoid(F.linear(x[..., :12], attn_gate_w)).view(B, T, self.num_heads, 1)
        y = y.contiguous().view(B, T, self.num_heads * self.head_dim)
        y = F.linear(y, sa_lambdas[1] * qkvo_w[self.dim * 3:].type_as(y))
        return y

class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        # Weights are stored in parameter banks and passed via forward()

    def forward(self, x: Tensor, c_fc: Tensor, c_proj: Tensor):
        # relu(x)^2:
        # https://arxiv.org/abs/2109.08668v2; ~1-2% better than GELU; suggested by @SKYLINEZ007 and @Grad62304977
        # Fused triton kernel for relu(x @ W1.T)^2 @ W2.T
        return FusedLinearReLUSquareFunction.apply(x, c_fc, c_proj)

class Block(nn.Module):
    def __init__(self, dim: int, head_dim: int, num_heads: int, has_attn: bool, has_mlp: bool, use_paired_head: bool):
        super().__init__()
        # skip attention of blocks.6 (the 7th layer) by @YouJiacheng
        if has_attn:
            if use_paired_head:
                self.attn = PairedHeadCausalSelfAttention(dim, head_dim, num_heads)
            else:
                self.attn = CausalSelfAttention(dim, head_dim, num_heads)
        else:
            self.attn = None
        # skip MLP blocks for first MLP layer by @EmelyanenkoK
        self.mlp = MLP() if has_mlp else None

    def forward(self, x: Tensor, attn_args: AttnArgs, qkvo_w: Tensor = None, c_fc: Tensor = None, c_proj: Tensor = None):
        if self.attn is not None:
            x = x + self.attn(norm(x), attn_args, qkvo_w)
        if self.mlp is not None:
            x = x + self.mlp(norm(x), c_fc, c_proj)
        return x

# -----------------------------------------------------------------------------
# The main model

def next_multiple_of_n(v: float | int, *, n: int):
    return next(x for x in range(n, int(v) + 1 + n, n) if x >= v)

@dataclass
class ForwardScheduleConfig:
    mtp_weights: torch.Tensor
    ws_short: int
    ws_long: int

class GPT(nn.Module):
    def __init__(self, vocab_size: int, num_layers: int, num_heads: int, head_dim: int, model_dim: int, max_seq_len: int):
        super().__init__()
        self.num_layers = num_layers
        vocab_size = next_multiple_of_n(vocab_size, n=128)

        self.smear_gate = nn.Linear(12, 1, bias=False)
        nn.init.zeros_(self.smear_gate.weight)
        self.smear_gate.weight.label = 'smear_gate'

        self.skip_gate = nn.Linear(12, 1, bias=False)
        nn.init.zeros_(self.skip_gate.weight)
        self.skip_gate.weight.label = 'skip_gate'

        # token value embeddings by @KoszarskyB - inspired by @Grad62304977's value residual implementation following https://arxiv.org/abs/2410.17897
        # value embedding code simplification inspired by @ragulpr https://github.com/KellerJordan/modded-nanogpt/pull/78
        self.value_embeds = nn.ModuleList([nn.Embedding(vocab_size, model_dim) for _ in range(5)])
        for embed in self.value_embeds:
            nn.init.zeros_(embed.weight)
        for i, ve in enumerate(self.value_embeds):
            ve.weight.label = f've{i}'  # ve0, ve1, ve2, ve3, ve4
        
        # parameter banks for attention and value embedding gate weights
        self.attn_gate_bank = nn.Parameter(torch.zeros(10, num_heads, 12)) # 10 layers
        self.attn_gate_bank.label = 'attn_gate_bank'
        self.ve_gate_bank = nn.Parameter(torch.zeros(5, num_heads, 12)) # 5 unique gates
        self.ve_gate_bank.label = 've_gate_bank'

        # -----------------------------------
        # Parameter banks for sharded optimization, by @chrisjmccormick

        # Identify which layers have attention/MLP
        # Attention is skipped in layer 6 by @YouJiacheng
        self.attn_layer_indices = [i for i in range(num_layers) if i != 6]
        # All layers have MLP (At 11 layers--dropped first layer @EmelyanenkoK)
        self.mlp_layer_indices = list(range(num_layers))

        hdim = num_heads * head_dim
        mlp_hdim = 4 * model_dim

        # Create index mappings: layer_idx -> bank_idx
        self.layer_to_attn_idx = {layer_idx: bank_idx for bank_idx, layer_idx in enumerate(self.attn_layer_indices)}
        self.layer_to_mlp_idx = {layer_idx: bank_idx for bank_idx, layer_idx in enumerate(self.mlp_layer_indices)}

        # Attention bank: stores QKVO weights for all attention layers
        # merged QKVO weights: suggested by many, implemented by @fernbear.bsky.social, and further improved by @YouJiacheng
        # https://x.com/hi_tysam/status/1879699187107033311
        # Simplified layout by @chrisjmccormick
        # Shape: (num_attn_layers, 4*model_dim, hdim) = (10, 3072, 768)
        # Reshape for sharding: (40, 768, 768) for even distribution across 8 GPUs
        self.attn_bank = nn.Parameter(torch.empty(len(self.attn_layer_indices), 4 * model_dim, hdim))
        self.attn_bank.label = 'attn'
        self.attn_bank.reshape = (len(self.attn_layer_indices) * 4, hdim, hdim)  # (40, 768, 768)

        # MLP bank: stores c_fc and c_proj for all MLP layers
        # Shape: (num_mlp_layers + padding, 2, mlp_hdim, model_dim) = (12, 2, 3072, 768)
        # We add 1 padding layer (index 11) to get 12*2=24 matrices for even distribution across 8 GPUs
        # Reshape for sharding: (24, 3072, 768)
        num_mlp_with_padding = len(self.mlp_layer_indices) + 1  # 11 + 1 = 12
        self.mlp_bank = nn.Parameter(torch.empty(num_mlp_with_padding, 2, mlp_hdim, model_dim))
        self.mlp_bank.label = 'mlp'
        self.mlp_bank.reshape = (num_mlp_with_padding * 2, mlp_hdim, model_dim)  # (24, 3072, 768)

        # improved init scale by @YouJiacheng
        # Attention uses dim^-0.5, MLP uses 0.5 * dim^-0.5
        attn_std = model_dim ** -0.5
        attn_bound = (3 ** 0.5) * attn_std
        mlp_std = 0.5 * (model_dim ** -0.5)
        mlp_bound = (3 ** 0.5) * mlp_std
        with torch.no_grad():
            # Init attention bank (QKV uniform, O zero)
            self.attn_bank[:, :model_dim * 3, :].uniform_(-attn_bound, attn_bound)
            self.attn_bank[:, model_dim * 3:, :].zero_()
            # Init MLP bank (c_fc uniform, c_proj zero) 
            self.mlp_bank[:, 0, :, :].uniform_(-mlp_bound, mlp_bound)  # c_fc
            self.mlp_bank[:, 1, :, :].zero_()  # c_proj - zero init suggested by @Grad62304977

        # Create blocks with has_attn/has_mlp flags
        self.paired_head_layers = [0, 2, 5, 9]
        self.blocks = nn.ModuleList([
            Block(model_dim, head_dim, num_heads, 
                  has_attn=(i in self.layer_to_attn_idx), 
                  has_mlp=(i in self.layer_to_mlp_idx),
                  use_paired_head=(i in self.paired_head_layers))
            for i in range(num_layers)
        ])
        self.yarn = Yarn(head_dim, max_seq_len)
        self.yarn_paired_head = YarnPairedHead(head_dim, max_seq_len)
        # there are only 50257 unique GPT-2 tokens; we extend to nearest multiple of 128 for efficiency.
        # suggested to me by @Grad62304977. this originates from Karpathy's experiments.
        use_fp8 = not os.environ.get("DISABLE_FP8", False)
        # Transposed weight storage for faster gradient accumulation
        self.lm_head = CastedLinearT(model_dim, vocab_size, use_fp8=use_fp8, x_s=100/448, w_s=1.6/448, grad_s=0.75/448)

        nn.init.normal_(self.lm_head.weight, mean=0, std=0.005)
        self.lm_head.weight.label = 'lm_head'

        self.embed = nn.Embedding(vocab_size, model_dim)
        self.embed.weight.label = 'embed'
        with torch.no_grad():
            self.embed.weight.copy_(self.lm_head.weight.T)

        self.bigram_embed = nn.Embedding(args.bigram_vocab_size, model_dim)
        self.bigram_embed.weight.label = 'bigram_embed'
        nn.init.zeros_(self.bigram_embed.weight)

        # x0_lambdas separated out for different optimizer treatment (no beta smoothing)
        self.x0_lambdas = nn.Parameter(torch.zeros(num_layers))
        self.x0_lambdas.label = 'x0_lambdas'

        pad = (-num_layers * 3 - 3) % dist.get_world_size()  # updated: 3*num_layers instead of 4*
        self.scalars = nn.Parameter(
            torch.cat(
                [
                    1.1 * torch.ones(num_layers),  # resid lambdas. 1.1 init such that layer i weight is i^(num_layers-i).
                    *[torch.tensor([0.5, 1.0]) for _ in range(num_layers)],  # SA lambdas
                    0.1 * torch.ones(num_layers), # bigram lambdas
                    torch.zeros(1), # smear_lambda
                    0.5*torch.ones(1), # backout_lambda
                    -1.5 * torch.ones(1),  # skip_lambda -> σ(-1.5) ≈ 0.18
                    torch.ones(pad),
                ]
            )
        )
        self.scalars.label = 'scalars'

    def forward(self, input_seq: Tensor, target_seq: Tensor, seqlens: Tensor, bigram_input_seq: Tensor, schedule_cfg: ForwardScheduleConfig):
        assert input_seq.ndim == 1

        # unpack schedule_cfg
        mtp_weights, ws_short, ws_long = schedule_cfg.mtp_weights, schedule_cfg.ws_short, schedule_cfg.ws_long

        # set configs
        skip_connections = []
        skip_in = [3] # long attention window on layer 3
        skip_out = [6] # no attn op on layer 6
        x_backout = None
        backout_layer = 7

        # set lambdas
        resid_lambdas = self.scalars[: 1 * self.num_layers]
        x0_lambdas = self.x0_lambdas
        sa_lambdas = self.scalars[1 * self.num_layers: 3 * self.num_layers].view(-1, 2)
        bigram_lambdas = self.scalars[3 * self.num_layers: 4 * self.num_layers]
        smear_lambda = self.scalars[4 * self.num_layers]
        backout_lambda = self.scalars[4 * self.num_layers+1]
        skip_lambda = self.scalars[4 * self.num_layers+2]

        # set block masks and key shift
        short_bm = ws_short * args.block_size
        long_bm = ws_long * args.block_size
        bm_sizes = [short_bm, short_bm, short_bm, long_bm, short_bm, short_bm, None, short_bm, short_bm, short_bm, long_bm]
        assert len(bm_sizes) == self.num_layers
        key_offset = [b==long_bm for b in bm_sizes] # apply partial key offset to long windows

        # Embedding lookup - embed is synced from lm_head during tied phase by optimizer
        x = self.embed(input_seq)
        x0_bigram = self.bigram_embed(bigram_input_seq)[None]
        
        # Value embeddings - always computed (not precomputed)
        ve = [value_embed(input_seq) for value_embed in self.value_embeds]
        # 01 ... 01 structure on token value embeddings by @YouJiacheng, improved on @leloykun's U-net structure
        # shifting first layer updates this to 01 ... 01 @photomz
        ve = [ve[0], ve[1]] + [None] * (self.num_layers - 5) + [ve[2], ve[3], ve[4]]
        assert len(ve) == self.num_layers

        # smear token embed forward 1 position @classiclarryd
        smear_gate_out = smear_lambda * torch.sigmoid(self.smear_gate(x[1:, :self.smear_gate.weight.size(-1)]))
        x = torch.cat([x[:1], x[1:] + smear_gate_out * x[:-1]])
        x = x0 = norm(x[None])

        # unbind gate banks to avoid select_backwards kernel
        ag = [w.bfloat16() for w in self.attn_gate_bank.unbind(0)] 
        veg = [w.bfloat16() for w in self.ve_gate_bank.unbind(0)]
        attn_gates = ag[:6] + [None] + ag[6:]
        ve_gates = [veg[0], veg[1]] + [None] * (self.num_layers - 5) + [veg[2], veg[3], veg[4]]
        assert len(attn_gates) == self.num_layers
        assert len(ve_gates) == self.num_layers

        # unbind weight banks to avoid select_backwards kernel
        attn_weights = self.attn_bank.unbind(0)  # tuple of [4*dim, hdim] tensors
        mlp_fcs = self.mlp_bank[:, 0, :, :].unbind(0)  # tuple of [mlp_hdim, dim] tensors
        mlp_projs = self.mlp_bank[:, 1, :, :].unbind(0)  # tuple of [mlp_hdim, dim] tensors

        for i in range(self.num_layers):
            yarn = self.yarn_paired_head if i in self.paired_head_layers else self.yarn
            attn_args = AttnArgs(
                ve=ve[i],
                sa_lambdas=sa_lambdas[i],
                seqlens=seqlens,
                bm_size=bm_sizes[i],
                yarn=yarn,
                key_offset=key_offset[i],
                attn_gate_w=attn_gates[i],
                ve_gate_w=ve_gates[i]
            )
            if i in skip_out:
                skip_gate_out = torch.sigmoid(skip_lambda) * 2 * torch.sigmoid(self.skip_gate(x0[..., :self.skip_gate.weight.size(-1)]))
                x = x + skip_gate_out * skip_connections.pop()
            if i == 0:
                x = (resid_lambdas[0] + x0_lambdas[0]) * x + bigram_lambdas[0] * x0_bigram
            else:
                x = resid_lambdas[i] * x + x0_lambdas[i] * x0 + bigram_lambdas[i] * x0_bigram
            
            # Get weights for this layer from banks
            qkvo_w = attn_weights[self.layer_to_attn_idx[i]] if i in self.layer_to_attn_idx else None
            c_fc = mlp_fcs[self.layer_to_mlp_idx[i]] if i in self.layer_to_mlp_idx else None
            c_proj = mlp_projs[self.layer_to_mlp_idx[i]] if i in self.layer_to_mlp_idx else None
            
            x = self.blocks[i](x, attn_args, qkvo_w, c_fc, c_proj)
            if i in skip_in:
                skip_connections.append(x)
            if i == backout_layer:
                x_backout = x

        # back out contributions from first 7 layers that are only required for downstream context and not direct prediction
        x -= backout_lambda * x_backout
        x = norm(x)
        logits = self.lm_head(x)
        # @Grad62304977 added tanh softcapping following Gemma 2 paper, @KoszarskyB reduced it from 30 to 15
        # @YouJiacheng shifted it by +15 (2*sigmoid(2*x)=tanh(x)+1). @classiclarryd updated to 23*sigmoid((logits+5)/7.5)
        if self.training:
            losses = FusedSoftcappedCrossEntropy.apply(logits.view(-1, logits.size(-1)), target_seq, mtp_weights)
            loss = losses.sum()
        else:
            logits = 23 * torch.sigmoid((logits + 5) / 7.5)
            logits_for_loss = logits.float()
            loss = F.cross_entropy(logits_for_loss.view(-1, logits_for_loss.size(-1)), target_seq, reduction="mean")
        return loss
# -----------------------------------------------------------------------------
# Distributed data loader

def _load_data_shard(file: Path):
    header = torch.from_file(str(file), False, 256, dtype=torch.int32) # header is 256 int32
    assert header[0] == 20240520, "magic number mismatch in the data .bin file"
    assert header[1] == 1, "unsupported version"
    num_tokens = int(header[2]) # number of tokens (claimed)
    with file.open("rb", buffering=0) as f:
        tokens = torch.empty(num_tokens, dtype=torch.uint16, pin_memory=True) # avoid pin_memory copy by @YouJiacheng
        f.seek(256 * 4)
        nbytes = f.readinto(tokens.numpy()) # avoid bytes->array copy by @YouJiacheng
        assert nbytes == 2 * num_tokens, "number of tokens read does not match header"
    return tokens

BOS_ID = 50256

class BOSFinder:
    # Helper for getting sequences that start at the beginning of documents by @varunneal based on work by @classiclarryd
    def __init__(self, tokens: Tensor, world_size: int = 1, quickload: bool = False):
        # Precompute BOS positions once per shard
        self.tokens=tokens
        self.size = tokens.numel()
        self.quickload = quickload
        if quickload:
            # only scan first 4 million tokens, then kickoff async thread to scan rest
            self.bos_idx = (tokens[:4_000_000] == BOS_ID).nonzero(as_tuple=True)[0].to(torch.int64).cpu().numpy()
            self.thread = None
            self.ready = threading.Event()
            self.start()
        else:
            self.bos_idx = (tokens == BOS_ID).nonzero(as_tuple=True)[0].to(torch.int64).cpu().numpy()
        self.i = 0
        self.world_size = world_size
        self.batch_iter = 0

    def _load(self):
        self.bos_idx_async = (self.tokens == BOS_ID).nonzero(as_tuple=True)[0].to(torch.int64).cpu().numpy()
        self.ready.set()

    def start(self):
        self.ready.clear()
        self.thread = threading.Thread(target=self._load)
        self.thread.start()

    def get(self):
        if self.thread:
            self.ready.wait()
            self.thread.join()
        self.bos_idx = self.bos_idx_async

    def next_batch(self, num_tokens_local: int, max_seq_len: int):
        # if quickload was used, repoint to the full dataset after 5 batches
        if self.quickload and self.batch_iter==5:
            self.get()
        n = len(self.bos_idx)
        starts = [[] for _ in range(self.world_size)]
        ends = [[] for _ in range(self.world_size)]

        idx = self.i
        for r in range(self.world_size):
            cur_len = 0
            while cur_len <= num_tokens_local:
                if idx >= n:
                    raise StopIteration(f"Insufficient BOS ahead; hit tail of shard.")
                cur = self.bos_idx[idx]
                starts[r].append(cur)
                end = min(self.bos_idx[idx + 1] if idx + 1 < n else self.size,
                          cur + max_seq_len,
                          cur + num_tokens_local - cur_len + 1)
                ends[r].append(end)
                cur_len += end - cur
                idx += 1

            assert cur_len == num_tokens_local + 1
        self.i = idx
        self.batch_iter+=1
        return starts, ends

class DataPreloader:
    # Helper for asynchronously loading next shard and indexing bos tokens
    def __init__(self, file_iter, world_size: int = 1):
        self.file_iter = file_iter
        self.world_size = world_size
        self.thread = None
        self.data = None
        self.ready = threading.Event()

    def _load(self):
        tokens = _load_data_shard(next(self.file_iter))
        self.data = (tokens, BOSFinder(tokens, self.world_size))
        self.ready.set()

    def start(self):
        self.ready.clear()
        self.thread = threading.Thread(target=self._load)
        self.thread.start()

    def get(self):
        if self.thread:
            self.ready.wait()
            self.thread.join()
        return self.data

def get_bigram_hash(x):
    """
    Computes bigram hash for each position using [prev_token, curr_token].
    Multiply by arbitary large ints to get even spread over int32 range.
    Position 0 is mapped to the reserved index (vocab_size - 1).
    BOS_tokens within the batch will hash based on last token of prior doc. Masking this ran slower and showed no improvement.
    """
    rand_int_1 = 36313
    rand_int_2 = 27191
    mod = args.bigram_vocab_size-1
    x = x.to(torch.int32).clone()
    x[0] = mod
    x[1:] = torch.bitwise_xor(rand_int_1 * x[1:], rand_int_2 * x[:-1]) % mod
    return x

def distributed_data_generator(filename_pattern: str, num_tokens: int, max_seq_len: int, grad_accum_steps: int = 1, align_to_bos: bool = True):
    # align_to_bos: each sequence begins with Beginning of Sequence token, sequences truncated to max_seq_len
    rank = dist.get_rank() if dist.is_initialized() else 0
    world_size = dist.get_world_size() if dist.is_initialized() else 1
    assert num_tokens % (world_size * grad_accum_steps) == 0, "Batch size must be divisible by world size"
    num_tokens = num_tokens // grad_accum_steps

    files = [Path(file) for file in sorted(glob.glob(filename_pattern))]
    if not files:
        raise FileNotFoundError(f"No files found for pattern: {filename_pattern}")

    file_iter = iter(files)  # Use itertools.cycle(files) for multi-epoch training
    tokens = _load_data_shard(next(file_iter))
    if align_to_bos:
        finder = BOSFinder(tokens, world_size=world_size, quickload=True)
        preloader = DataPreloader(file_iter, world_size)
        preloader.start()
    else:
        pos = 0  # for unaligned case

    while True:
        num_tokens_local = num_tokens // world_size
        max_num_docs = next_multiple_of_n(num_tokens_local // 300, n=128)  # median doc length is ~400

        if align_to_bos:
            try:
                seq_starts, seq_ends = finder.next_batch(num_tokens_local, max_seq_len)
                start_idxs, end_idxs = torch.tensor(seq_starts[rank]), torch.tensor(seq_ends[rank])
            except StopIteration:
                # This shard is exhausted, load the next one in the next loop iteration.
                tokens, finder = preloader.get()
                preloader.start()
                continue

            buf = torch.cat([tokens[i:j] for i, j in zip(start_idxs, end_idxs)])
            _inputs = buf[:-1]
            _targets = buf[1:]
            end_idxs[-1] -= 1  # last document was too long to account for _targets offset
            cum_lengths = (end_idxs - start_idxs).cumsum(0)

        else:
            if pos + num_tokens + 1 >= len(tokens):  # should not occur for val data
                tokens, pos = _load_data_shard(next(file_iter)), 0

            pos_local = pos + rank * num_tokens_local
            buf = tokens[pos_local: pos_local + num_tokens_local + 1]
            _inputs = buf[:-1].view(num_tokens_local, )
            _targets = buf[1:].view(num_tokens_local, )

            cum_lengths = torch.nonzero(_inputs == BOS_ID)[:, 0]
            pos += num_tokens


        _cum_lengths = torch.full((max_num_docs,), num_tokens_local)
        _cum_lengths[0] = 0
        _cum_lengths[1:len(cum_lengths) + 1] = cum_lengths

        # Cast to int32 on CPU before transfer to avoid dtype conversion during .to()
        _inputs = _inputs.to(dtype=torch.int32)
        _targets = _targets.to(dtype=torch.int64)
        _cum_lengths = _cum_lengths.to(dtype=torch.int32)
        _bigram_inputs = get_bigram_hash(_inputs)

        new_params = yield (
            _inputs.to(device="cuda", non_blocking=True),
            _targets.to(device="cuda", non_blocking=True),
            _cum_lengths.to(device="cuda", non_blocking=True),
            _bigram_inputs.to(device="cuda", non_blocking=True)
        )

        if new_params is not None:
            # makes it possible for generator to receive new (num_tokens, max_seq_len, grad_accum_steps) via .send()
            new_num_tokens, new_max_seq_len, new_grad_accum_steps = new_params
            assert new_num_tokens % (world_size * new_grad_accum_steps) == 0, "Num tokens must be divisible by world size"
            num_tokens = new_num_tokens // new_grad_accum_steps
            max_seq_len = new_max_seq_len

# -----------------------------------------------------------------------------
# Training Management
 
def get_bs(step: int):
    if step >= args.num_scheduled_iterations:
        return args.train_bs_extension
    x = step / args.num_scheduled_iterations
    bs_idx = int(len(args.train_bs_schedule) * x)
    return args.train_bs_schedule[bs_idx]

def get_ws(step: int):
    # set short window size to half of long window size
    # Higher ws on "extension" steps
    if step >= args.num_scheduled_iterations:
        return args.ws_final // 2, args.ws_final
    x = step / args.num_scheduled_iterations
    assert 0 <= x < 1
    ws_idx = int(len(args.ws_schedule) * x)
    return args.ws_schedule[ws_idx] // 2, args.ws_schedule[ws_idx]

# learning rate schedule: tied to batch size schedule, with cooldown at the end.
def get_lr(step: int):
    if step > args.num_scheduled_iterations:
        return 0.1
    lr_max = 1.0
    x = step / args.num_scheduled_iterations
    if x > 1/3:
       lr_max = 1.52  # (16/8)**0.6
    if x > 2/3:
        lr_max = 1.73  # (24/8)**0.5
    if x >= 1 - args.cooldown_frac:
        w = (1 - x) / args.cooldown_frac
        lr = lr_max * w + (1 - w) * 0.1
        return lr
    return lr_max

def get_muon_momentum(step: int, muon_warmup_steps=300, muon_cooldown_steps=50, momentum_min=0.85, momentum_max=0.95):
    # warmup phase: linearly increase momentum from min to max
    # cooldown phase: linearly decrease momentum from max to min
    momentum_cd_start = args.num_iterations - muon_cooldown_steps
    if step < muon_warmup_steps:
        frac = step / muon_warmup_steps
        momentum = momentum_min + frac * (momentum_max - momentum_min)
    elif step > momentum_cd_start:
        frac = (step - momentum_cd_start) / muon_cooldown_steps
        momentum = momentum_max - frac * (momentum_max - momentum_min)
    else:
        momentum = momentum_max
    return momentum

class TrainingManager():
    """
    Manages the NorMuonAndAdam for all parameters with explicit ordering.
    Notable Features:
        1. Scalars are given higher momentum terms to smooth learning @ChrisJMcCormick
        2. Adam optimizers are only stepped on odd steps @classiclarryd
        3. Explicit scatter_order and work_order for communication scheduling (no backward hooks)
        4. Muon has a linear momentum warmup and cooldown schedule
        5. Learning rates follow a linear decay schedule
        6. Embed is tied to lm_head until split step (2/3 of training), then untied @classiclarryd

    Manages model architecture, data, and target that changes during training
    Notable Features:
        1. Multi Token Prediction schedule of [1, 0.5, 0.25->0] -> [1, 0.5->0] -> [1] @varunneal
        2. Sliding Attention window schedule of [1,3] -> [3,7] -> [5,11] -> [6,13]
        3. YaRN updates to RoPE on window changes
        4. Split embed and lm_head at 2/3 of training (weights and optimizer state copied)
        5. Batch size schedule of 8 -> 16 -> 24
        6. Post training extension of long windows from 13 to 20
    """
    def __init__(self, model):
        self.mtp_weights_schedule = self._build_mtp_schedule()
        self.model = model
        
        # - Ordering dictates when to launch reduce/reduce_scatter operations
        # - "sharded" parameters use reduce_scatter/all_gather and "replicated" ones use all_reduce
        # - lr_mul and wd_mul are per-parameter learning rate and weight decay multipliers
        self.param_table = {
            "attn":           {"optim": "normuon", "comms": "sharded",    "adam_betas": None},
            "mlp":            {"optim": "normuon", "comms": "sharded",    "adam_betas": None},         
            "scalars":        {"optim": "adam",    "comms": "replicated", "adam_betas": [0.9,  0.99], "lr_mul": 5.0,  "wd_mul": 0.0},
            "ve0":            {"optim": "adam",    "comms": "sharded",    "adam_betas": [0.75, 0.95], "lr_mul": 75.,  "wd_mul": 5.0},
            "ve1":            {"optim": "adam",    "comms": "sharded",    "adam_betas": [0.75, 0.95], "lr_mul": 75.,  "wd_mul": 5.0},
            "ve2":            {"optim": "adam",    "comms": "sharded",    "adam_betas": [0.75, 0.95], "lr_mul": 75.,  "wd_mul": 5.0},
            "ve3":            {"optim": "adam",    "comms": "sharded",    "adam_betas": [0.75, 0.95], "lr_mul": 75.,  "wd_mul": 5.0},
            "ve4":            {"optim": "adam",    "comms": "sharded",    "adam_betas": [0.75, 0.95], "lr_mul": 75.,  "wd_mul": 5.0},
            "bigram_embed":   {"optim": "adam",    "comms": "sharded",    "adam_betas": [0.75, 0.95], "lr_mul": 75.,  "wd_mul": 5.0},
            "smear_gate":     {"optim": "adam",    "comms": "replicated", "adam_betas": [0.9,  0.99], "lr_mul": 0.01, "wd_mul": 0.0},
            "skip_gate":      {"optim": "adam",    "comms": "replicated", "adam_betas": [0.9,  0.99], "lr_mul": 0.05, "wd_mul": 0.0},
            "attn_gate_bank": {"optim": "adam",    "comms": "replicated", "adam_betas": [0.9,  0.99]},
            "ve_gate_bank":   {"optim": "adam",    "comms": "replicated", "adam_betas": [0.9,  0.99]},
            "x0_lambdas":     {"optim": "adam",    "comms": "replicated", "adam_betas": [0.65, 0.95], "lr_mul": 5.0,  "wd_mul": 0.0},
            "lm_head":        {"optim": "adam",    "comms": "sharded",    "adam_betas": [0.5,  0.95], "wd_mul": 150.},
            "embed":          {"optim": "adam",    "comms": "sharded",    "adam_betas": [0.5,  0.95], "wd_mul": 150.},
        }

        # - Process smaller/faster params first while large reduces complete
        # - lm_head must complete before embed sync (when tied)
        self.work_order = [
            "scalars", "smear_gate", "skip_gate", "attn_gate_bank", "ve_gate_bank", "x0_lambdas",  # Small, fast
            "ve0", "ve1", "ve2", "ve3", "ve4", "bigram_embed",  # Medium
            "lm_head", "embed",   # lm_head must complete before embed sync (when tied)
            "attn", "mlp",        # Large, polar express - process last to maximize overlap
        ]

        adam_defaults = dict(
            lr=0.008,
            eps=1e-10,
            weight_decay=0.005,
        )
        
        normuon_defaults = dict(
            lr=0.023,
            momentum=0.95,
            beta2=0.95,
            weight_decay=1.2,
        )
        
        self.optimizer = NorMuonAndAdam(
            model.named_parameters(),
            param_table=self.param_table,
            scatter_order=list(self.param_table.keys()),  # Dict order defines scatter priority
            work_order=self.work_order,
            adam_defaults=adam_defaults,
            normuon_defaults=normuon_defaults,
        )

        # Split embed from lm_head at 2/3 of training (on an odd step so Adam updates)
        self.split_step = math.ceil(args.split_embed_frac * args.num_scheduled_iterations) | 1

        self.reset()

    def _build_mtp_schedule(self):
        # Precompute MTP weights for all steps to avoid tensor allocation during training
        # Schedule: [1, 0.5, 0.25->0] -> [1, 0.5->0] -> [1]
        mtp_weights_schedule = []
        for s in range(args.num_iterations + 1):
            x = s / args.num_scheduled_iterations
            if x < 1/3:
                w = [1.0, 0.5, 0.25 * (1 - 3*x)]
            elif x < 2/3:
                w = [1.0, 0.5 * (1 - (3*x - 1))]
            else:
                w = [1.0]
            mtp_weights_schedule.append(torch.tensor(w, device=device))
        return mtp_weights_schedule

    def apply_final_ws_ext(self):
        self.ws_long = args.ws_validate_post_yarn_ext

    def get_forward_args(self):
        return ForwardScheduleConfig(
            mtp_weights = self.mtp_weights,
            ws_short = self.ws_short,
            ws_long = self.ws_long
        )
    
    def _is_adam_step(self, step: int):
        """Adam params are only updated on odd steps."""
        return step % 2 == 1

    def get_transition_steps(self):
        transition_steps = []
        ws_short, ws_long = get_ws(0)
        for step in range(1, args.num_iterations):
            ws_short, new_ws_long = get_ws(step)
            if new_ws_long != ws_long:
                transition_steps.append(step)
                ws_long = new_ws_long
        return transition_steps

    def advance_schedule(self, step: int):
        self.ws_short, new_ws_long = get_ws(step)
        if new_ws_long != self.ws_long:
            self.model.yarn.apply(self.ws_long, new_ws_long)
            self.model.yarn_paired_head.apply(self.ws_long, new_ws_long)

        new_batch_size = get_bs(step)
        if new_batch_size != self.batch_size:
            self.train_loader_send_args = (new_batch_size, args.train_max_seq_len, grad_accum_steps)
            self.batch_size = new_batch_size
        else:
            self.train_loader_send_args = None

        self.ws_long = new_ws_long
        self.mtp_weights = self.mtp_weights_schedule[step]
    
    def step_optimizers(self, step: int):
        step_lr = get_lr(step)
        muon_momentum = get_muon_momentum(step)
        do_adam = self._is_adam_step(step)
        
        # Update learning rates and momentum for all params
        for param, p_cfg in self.optimizer.param_cfgs.items():
            p_cfg.lr = p_cfg.initial_lr * step_lr
            if p_cfg.optim == "normuon":
                p_cfg.momentum = muon_momentum
        
        # Step optimizer with do_adam flag
        self.optimizer.step(do_adam=do_adam)
        
        # At split step: copy lm_head optimizer state to embed and mark as split
        if step == self.split_step:
            self.optimizer.copy_lm_state_to_embed()

    def reset(self, state=None):
        if state is not None:
            self.optimizer.load_state_dict(state)

        # Reset NorMuon momentum buffers and split_embed state
        self.optimizer.reset()

        self.ws_short, self.ws_long = get_ws(0)
        self.batch_size = get_bs(0)
        self.model.yarn.reset()
        self.model.yarn_paired_head.reset()

    def get_state(self):
        return copy.deepcopy(self.optimizer.state_dict())

# -----------------------------------------------------------------------------
# int main

@dataclass
class Hyperparameters:
    # data
    train_files: str = "data/fineweb10B/fineweb_train_*.bin" # input .bin to train on
    val_files: str = "data/fineweb10B/fineweb_val_*.bin" # input .bin to eval validation loss on
    val_tokens: int = 10485760 # how many tokens of validation data? it's important to keep this fixed for consistent comparisons
    # batch sizes
    train_bs_schedule: tuple = (8 * 2048 * 8, 16 * 2048 * 8, 24 * 2048 * 8)
    train_bs_extension: int = 24 * 2048 * 8
    train_max_seq_len: int = 128 * 16
    val_batch_size: int = 4 * 64 * 1024 * 8
    # optimization
    num_scheduled_iterations: int = 1535  # number of steps to complete lr and ws schedule
    num_extension_iterations: int = 40  # number of steps to continue training at final lr and ws
    num_iterations: int = num_scheduled_iterations + num_extension_iterations
    cooldown_frac: float = 0.55  # fraction of num_scheduled_iterations spent cooling down the learning rate
    split_embed_frac: float = 2/3  # fraction of training when embeddings split from lm_head
    # evaluation and logging
    run_id: str = f"{uuid.uuid4()}"
    val_loss_every: int = 250  # every how many steps to evaluate val loss? 0 for only at the end
    save_checkpoint: bool = False
    # attention masking
    block_size: int = 128
    ws_schedule: tuple = (3, 7, 11)
    ws_final: int = 13 # increase final validation ws, used for YaRN extension and short window size @classiclarryd
    ws_validate_post_yarn_ext: int = 20 # extend long windows out even further after applying YaRN
    # bigram hash embedding
    bigram_vocab_size = 50304 * 5

args = Hyperparameters()

data_path = os.environ.get("DATA_PATH", ".")
args.train_files = os.path.join(data_path, args.train_files)
args.val_files = os.path.join(data_path, args.val_files)

# torchrun sets these env variables
rank = int(os.environ["RANK"])
world_size = int(os.environ["WORLD_SIZE"])
assert 8 % world_size == 0, "world_size must be a divisor of 8"
grad_accum_steps = 8 // world_size
assert torch.cuda.is_available()
device = torch.device("cuda", int(os.environ["LOCAL_RANK"]))
torch.cuda.set_device(device)
dist.init_process_group(backend="nccl", device_id=device)
dist.barrier()
master_process = (rank == 0) # this process will do logging, checkpointing etc.

# begin logging
logfile = None
if master_process:
    run_id = args.run_id
    os.makedirs("logs", exist_ok=True)
    logfile = f"logs/{run_id}.txt"
    print(logfile)
def print0(s, console=False):
    if master_process:
        with open(logfile, "a") as f:
            if console:
                print(s)
            print(s, file=f)

# begin by printing this file (the Python code)
print0(code)
print0("="*100)
# log information about the hardware/software environment this is running on
print0(f"Running Python {sys.version}")
print0(f"Running PyTorch {torch.version.__version__} compiled for CUDA {torch.version.cuda}")
print0(f"Running Triton version {triton.__version__}")

def nvidia_smi():
    import subprocess  # avoid top level import
    return subprocess.run(["nvidia-smi"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True).stdout
print0(nvidia_smi())
print0("="*100)

model: nn.Module = GPT(
    vocab_size=50257,
    num_layers=11,
    num_heads=6,
    head_dim=128,
    model_dim=768,
    max_seq_len=args.val_batch_size // (grad_accum_steps * world_size)
).cuda()
for m in model.modules():
    if isinstance(m, (nn.Embedding, nn.Linear)):
        m.weight.data = m.weight.data.bfloat16()
model.attn_gate_bank.data = model.attn_gate_bank.data.bfloat16()
model.ve_gate_bank.data = model.ve_gate_bank.data.bfloat16()
model.attn_bank.data = model.attn_bank.data.bfloat16()
model.mlp_bank.data = model.mlp_bank.data.bfloat16()
for param in model.parameters():
    dist.broadcast(param.detach(), 0)

model: nn.Module = torch.compile(model, dynamic=False, fullgraph=True)
training_manager = TrainingManager(model)

########################################
#            Warmup kernels            #
########################################
print0("Compiling model and warming up kernels (~7 minutes on first execution)", console=True)
# Warmup the training kernels, then re-initialize the state so we aren't cheating
initial_state = dict(model=copy.deepcopy(model.state_dict()),
                     optimizer=training_manager.get_state()) # save the initial state
train_loader = distributed_data_generator(args.train_files, args.train_bs_schedule[0], args.train_max_seq_len, grad_accum_steps=grad_accum_steps)
val_loader = distributed_data_generator(args.val_files, args.val_batch_size, -1, grad_accum_steps=grad_accum_steps, align_to_bos=False)

transition_steps = training_manager.get_transition_steps()
# first few steps plus transitions
warmup_steps = sorted({0, 1, 2} | set(s + offset for s in transition_steps for offset in [-1, 0, 1] if s + offset >= 0)) 
print0(f"Sampling steps {warmup_steps} for warmup", console=True)
for step in warmup_steps:
    training_manager.advance_schedule(step)
    model.eval()
    with torch.no_grad():
        inputs, targets, cum_seqlens, bigram_inputs = next(val_loader)
        model(inputs, targets, cum_seqlens, bigram_inputs, training_manager.get_forward_args())
    model.train()
    for idx in range(grad_accum_steps):
        send_args = training_manager.train_loader_send_args
        inputs, targets, cum_seqlens, bigram_inputs = train_loader.send(send_args)
        (model(inputs, targets, cum_seqlens, bigram_inputs, training_manager.get_forward_args()) / grad_accum_steps).backward()
    training_manager.step_optimizers(step)
print0("Resetting Model", console=True)
model.zero_grad(set_to_none=True)
model.load_state_dict(initial_state["model"])
training_manager.reset(initial_state["optimizer"])
del val_loader, train_loader, initial_state
model.train()

########################################
#        Training and validation       #
########################################
train_loader = distributed_data_generator(args.train_files, args.train_bs_schedule[0], args.train_max_seq_len, grad_accum_steps=grad_accum_steps)

gc.collect()

training_time_ms = 0
# start the clock
torch.cuda.synchronize()
t0 = time.perf_counter()
# begin training
train_steps = args.num_iterations
for step in range(train_steps + 1):
    last_step = (step == train_steps)
    training_manager.advance_schedule(step)
    # --------------- VALIDATION SECTION -----------------
    if last_step or (args.val_loss_every > 0 and step % args.val_loss_every == 0):
        if last_step:
            training_manager.apply_final_ws_ext()
        # stop the clock
        torch.cuda.synchronize()
        training_time_ms += 1000 * (time.perf_counter() - t0)
        model.eval()
        assert args.val_tokens % args.val_batch_size == 0
        val_steps = grad_accum_steps * args.val_tokens // args.val_batch_size
        val_loader = distributed_data_generator(args.val_files, args.val_batch_size, -1, grad_accum_steps=grad_accum_steps, align_to_bos=False)
        val_loss = 0
        with torch.no_grad():
            for _ in range(val_steps):
                inputs, targets, cum_seqlens, bigram_inputs = next(val_loader)
                val_loss += model(inputs, targets, cum_seqlens, bigram_inputs, training_manager.get_forward_args())
        val_loss /= val_steps
        del val_loader
        dist.reduce(val_loss, 0, op=dist.ReduceOp.AVG)
        print0(f"step:{step}/{train_steps} val_loss:{val_loss:.4f} train_time:{training_time_ms:.0f}ms step_avg:{training_time_ms/max(step, 1):.2f}ms", console=True)
        model.train()
        # start the clock again
        torch.cuda.synchronize()
        t0 = time.perf_counter()

    if last_step:
        if master_process and args.save_checkpoint:
            log = dict(step=step, code=code, model=model.state_dict(), optimizer=training_manager.get_state())
            os.makedirs(f"logs/{run_id}", exist_ok=True)
            torch.save(log, f"logs/{run_id}/state_step{step:06d}.pt")
        # the last step only has the validation loop, so break to avoid training
        break

    # --------------- TRAINING SECTION -----------------
    for idx in range(grad_accum_steps):
        inputs, targets, cum_seqlens, bigram_inputs = train_loader.send(training_manager.train_loader_send_args)
        (model(inputs, targets, cum_seqlens, bigram_inputs, training_manager.get_forward_args()) / grad_accum_steps).backward()
    training_manager.step_optimizers(step)

    # logging
    approx_training_time_ms = training_time_ms + 1000 * (time.perf_counter() - t0)
    print0(f"step:{step+1}/{train_steps} train_time:{approx_training_time_ms:.0f}ms step_avg:{approx_training_time_ms/(step + 1):.2f}ms", console=True)

print0(f"peak memory allocated: {torch.cuda.max_memory_allocated() // 1024 // 1024} MiB "
       f"reserved: {torch.cuda.max_memory_reserved() // 1024 // 1024} MiB", console=True)
dist.destroy_process_group()

----------------------------------------
# triton_kernels.py
----------------------------------------

import torch
import triton
import triton.language as tl
from triton.tools.tensor_descriptor import TensorDescriptor

# -----------------------------------------------------------------------------
# Triton kernel for symmetric matrix multiplication by @byronxu99

def _get_autotune_configs():
    return [
        triton.Config(
            {
                "BLOCK_SIZE_M": bm,
                "BLOCK_SIZE_N": bn,
                "BLOCK_SIZE_K": bk,
                "GROUP_SIZE_M": 8,
                "LOWER_UPPER": 1,
            },
            num_stages=stages,
            num_warps=warps,
        )
        for bm in [64, 128]
        for bn in [64, 128, 256]
        for bk in [64, 128]
        for stages, warps in [(3, 4), (3, 8), (4, 4)]
        if bm // bn <= 2 and bn // bm <= 2
    ]

@triton.jit
def _pid_to_block(
    pid,
    M,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    GROUP_SIZE_M: tl.constexpr,
):
    # Split output matrix into blocks of size (BLOCK_SIZE_M, BLOCK_SIZE_N)
    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
    num_pid_n = tl.cdiv(M, BLOCK_SIZE_N)

    # Map PID to a single matrix in batch
    batch_idx = pid // (num_pid_m * num_pid_n)
    pid = pid % (num_pid_m * num_pid_n)

    # Map PID to 2D grid of blocks
    pid_m = pid // num_pid_n
    pid_n = pid % num_pid_n
    pid_m, pid_n = tl.swizzle2d(pid_m, pid_n, num_pid_m, num_pid_n, GROUP_SIZE_M)

    m_idx = pid_m * BLOCK_SIZE_M
    n_idx = pid_n * BLOCK_SIZE_N
    return batch_idx, m_idx, n_idx

@triton.autotune(
    configs=_get_autotune_configs(),
    key=["M", "K", "a_stride_r", "a_stride_c", "c_stride_r", "c_stride_c"],
)
@triton.jit
def XXT_kernel(
    A_ptr, C_ptr,
    M, K,
    a_stride_b, a_stride_r, a_stride_c,
    c_stride_b, c_stride_r, c_stride_c,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr,
    GROUP_SIZE_M: tl.constexpr,
    LOWER_UPPER: tl.constexpr,
):
    pid = tl.program_id(axis=0)
    batch_idx, m_idx, n_idx = _pid_to_block(
        pid, M, BLOCK_SIZE_M, BLOCK_SIZE_N, GROUP_SIZE_M
    )

    # Skip blocks that don't need to be computed
    skip_block_below_diag = (LOWER_UPPER == 0) and (n_idx + BLOCK_SIZE_N <= m_idx)
    skip_block_above_diag = (LOWER_UPPER != 0) and (m_idx + BLOCK_SIZE_M <= n_idx)
    if skip_block_below_diag or skip_block_above_diag:
        return

    # Index into one matrix of batch
    A_ptr += batch_idx * a_stride_b
    C_ptr += batch_idx * c_stride_b

    # Create pointer arrays for A and A.T
    offs_m = (m_idx + tl.arange(0, BLOCK_SIZE_M)) % M
    offs_n = (n_idx + tl.arange(0, BLOCK_SIZE_N)) % M
    offs_k = tl.arange(0, BLOCK_SIZE_K)
    a_ptrs = A_ptr + (offs_m[:, None] * a_stride_r + offs_k[None, :] * a_stride_c)
    at_ptrs = A_ptr + (offs_k[:, None] * a_stride_c + offs_n[None, :] * a_stride_r)

    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)

    # Accumulate over blocks of K
    for k in tl.range(0, tl.cdiv(K, BLOCK_SIZE_K)):
        a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0)
        at = tl.load(at_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
        accumulator = tl.dot(a, at, accumulator)
        a_ptrs += BLOCK_SIZE_K * a_stride_c
        at_ptrs += BLOCK_SIZE_K * a_stride_c

    out_dtype = C_ptr.dtype.element_ty
    output = accumulator.to(out_dtype)

    # Store block of C
    offs_cm = m_idx + tl.arange(0, BLOCK_SIZE_M)
    offs_cn = n_idx + tl.arange(0, BLOCK_SIZE_N)
    c_ptrs = C_ptr + (offs_cm[:, None] * c_stride_r + offs_cn[None, :] * c_stride_c)
    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < M)
    tl.store(c_ptrs, output, mask=c_mask)

    # Store block of C mirrored across the diagonal
    c_ptrs_t = C_ptr + (offs_cn[:, None] * c_stride_r + offs_cm[None, :] * c_stride_c)
    c_mask_t = (offs_cn[:, None] < M) & (offs_cm[None, :] < M)
    tl.store(c_ptrs_t, output.T, mask=c_mask_t)

def XXT(A: torch.Tensor, out: torch.Tensor):
    """
    Launch Triton kernel to compute C = A @ A.T
    """
    assert A.ndim == 2 or A.ndim == 3
    M, K = A.shape[-2:]
    assert out.size(-2) == M, "Output matrix has incorrect shape"
    assert out.size(-1) == M, "Output matrix has incorrect shape"

    batch_size = A.size(0) if A.ndim == 3 else 1
    input_batch_stride = A.stride(0) if A.ndim == 3 else 0
    output_batch_stride = out.stride(0) if out.ndim == 3 else 0

    grid = lambda meta: (
        batch_size * triton.cdiv(M, meta["BLOCK_SIZE_M"]) * triton.cdiv(M, meta["BLOCK_SIZE_N"]),
    )
    XXT_kernel[grid](
        A_ptr=A,
        C_ptr=out,
        M=M,
        K=K,
        a_stride_b=input_batch_stride,
        a_stride_r=A.stride(-2),
        a_stride_c=A.stride(-1),
        c_stride_b=output_batch_stride,
        c_stride_r=out.stride(-2),
        c_stride_c=out.stride(-1),
    )
    return out

@triton.autotune(
    configs=_get_autotune_configs(),
    key=["M", "a_stride_r", "a_stride_c", "c_stride_r", "c_stride_c"],
)
@triton.jit
def ba_plus_cAA_kernel(
    A_ptr, C_ptr,
    M,
    a_stride_b, a_stride_r, a_stride_c,
    c_stride_b, c_stride_r, c_stride_c,
    alpha, beta,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr,
    GROUP_SIZE_M: tl.constexpr,
    LOWER_UPPER: tl.constexpr,
):
    # This is mostly duplicated from XXT_kernel, but also loads and adds a block of A
    # Performance is slightly slower than XXT_kernel, so we use two separate kernels
    pid = tl.program_id(axis=0)
    batch_idx, m_idx, n_idx = _pid_to_block(
        pid, M, BLOCK_SIZE_M, BLOCK_SIZE_N, GROUP_SIZE_M
    )

    # Skip blocks that don't need to be computed
    skip_block_below_diag = (LOWER_UPPER == 0) and (n_idx + BLOCK_SIZE_N <= m_idx)
    skip_block_above_diag = (LOWER_UPPER != 0) and (m_idx + BLOCK_SIZE_M <= n_idx)
    if skip_block_below_diag or skip_block_above_diag:
        return

    # Index into one matrix of batch
    A_ptr += batch_idx * a_stride_b
    C_ptr += batch_idx * c_stride_b

    # Create pointer arrays for A and A.T
    offs_m = (m_idx + tl.arange(0, BLOCK_SIZE_M)) % M
    offs_n = (n_idx + tl.arange(0, BLOCK_SIZE_N)) % M
    offs_k = tl.arange(0, BLOCK_SIZE_K)
    a_ptrs = A_ptr + (offs_m[:, None] * a_stride_r + offs_k[None, :] * a_stride_c)
    at_ptrs = A_ptr + (offs_k[:, None] * a_stride_c + offs_n[None, :] * a_stride_r)

    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)

    # Accumulate over blocks of K
    for k in tl.range(0, tl.cdiv(M, BLOCK_SIZE_K)):
        a = tl.load(a_ptrs, mask=offs_k[None, :] < M - k * BLOCK_SIZE_K, other=0.0)
        at = tl.load(at_ptrs, mask=offs_k[:, None] < M - k * BLOCK_SIZE_K, other=0.0)
        accumulator = tl.dot(a, at, accumulator)
        a_ptrs += BLOCK_SIZE_K * a_stride_c
        at_ptrs += BLOCK_SIZE_K * a_stride_c

    # Load block of A to add (corresponds to the current block of C)
    offs_am = m_idx + tl.arange(0, BLOCK_SIZE_M)
    offs_an = n_idx + tl.arange(0, BLOCK_SIZE_N)
    a_add_ptrs = A_ptr + (offs_am[:, None] * a_stride_r + offs_an[None, :] * a_stride_c)
    a_add_mask = (offs_am[:, None] < M) & (offs_an[None, :] < M)
    a_add = tl.load(a_add_ptrs, mask=a_add_mask, other=0.0).to(tl.float32)

    # Apply alpha and beta
    accumulator *= alpha
    accumulator += a_add * beta

    out_dtype = C_ptr.dtype.element_ty
    output = accumulator.to(out_dtype)

    # Store block of C
    offs_cm = m_idx + tl.arange(0, BLOCK_SIZE_M)
    offs_cn = n_idx + tl.arange(0, BLOCK_SIZE_N)
    c_ptrs = C_ptr + (offs_cm[:, None] * c_stride_r + offs_cn[None, :] * c_stride_c)
    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < M)
    tl.store(c_ptrs, output, mask=c_mask)

    # Store block of C mirrored across the diagonal
    c_ptrs_t = C_ptr + (offs_cn[:, None] * c_stride_r + offs_cm[None, :] * c_stride_c)
    c_mask_t = (offs_cn[:, None] < M) & (offs_cm[None, :] < M)
    tl.store(c_ptrs_t, output.T, mask=c_mask_t)

def ba_plus_cAA(A: torch.Tensor, alpha: float, beta: float, out: torch.Tensor):
    """
    Launch Triton kernel to compute C = alpha * A @ A.T + beta * A
    """
    assert A.ndim == 2 or A.ndim == 3
    M, K = A.shape[-2:]
    assert M == K, "Input matrix must be square"
    assert out.size(-2) == M
    assert out.size(-1) == M

    batch_size = A.size(0) if A.ndim == 3 else 1
    input_batch_stride = A.stride(0) if A.ndim == 3 else 0
    output_batch_stride = out.stride(0) if out.ndim == 3 else 0

    grid = lambda meta: (
        batch_size * triton.cdiv(M, meta["BLOCK_SIZE_M"]) * triton.cdiv(M, meta["BLOCK_SIZE_N"]),
    )
    ba_plus_cAA_kernel[grid](
        A_ptr=A,
        C_ptr=out,
        M=M,
        a_stride_b=input_batch_stride,
        a_stride_r=A.stride(-2),
        a_stride_c=A.stride(-1),
        c_stride_b=output_batch_stride,
        c_stride_r=out.stride(-2),
        c_stride_c=out.stride(-1),
        alpha=alpha,
        beta=beta,
    )
    return out

# -----------------------------------------------------------------------------
# Triton kernel for MLP: relu(x @ W1.T)^2, by @andrewbriand, @jrauvola

@triton.jit
def linear_relu_square_kernel(a_desc, b_desc, c_desc, aux_desc,
                                 M, N, K,
                                 BLOCK_SIZE_M: tl.constexpr,
                                 BLOCK_SIZE_N: tl.constexpr,
                                 BLOCK_SIZE_K: tl.constexpr,
                                 GROUP_SIZE_M: tl.constexpr,
                                 NUM_SMS: tl.constexpr,
                                 FORWARD: tl.constexpr,
                                 ):
    dtype = tl.bfloat16
    start_pid = tl.program_id(axis=0)
    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
    k_tiles = tl.cdiv(K, BLOCK_SIZE_K)
    num_tiles = num_pid_m * num_pid_n

    tile_id_c = start_pid - NUM_SMS
    num_pid_in_group = GROUP_SIZE_M * num_pid_n

    for tile_id in tl.range(start_pid, num_tiles, NUM_SMS, flatten=True):
        pid_m = tile_id // num_pid_n
        pid_n = tile_id % num_pid_n
        offs_am = pid_m * BLOCK_SIZE_M
        offs_bn = pid_n * BLOCK_SIZE_N

        accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
        for ki in range(k_tiles):
            offs_k = ki * BLOCK_SIZE_K
            a = a_desc.load([offs_am, offs_k])
            b = b_desc.load([offs_bn, offs_k])
            accumulator = tl.dot(a, b.T, accumulator)

        tile_id_c += NUM_SMS
        pid_m = tile_id // num_pid_n
        pid_n = tile_id % num_pid_n
        offs_am_c = pid_m * BLOCK_SIZE_M
        offs_bn_c = pid_n * BLOCK_SIZE_N

        acc = tl.reshape(accumulator, (BLOCK_SIZE_M, 2, BLOCK_SIZE_N // 2))
        acc = tl.permute(acc, (0, 2, 1))
        acc0, acc1 = tl.split(acc)

        c0 = acc0.to(dtype)
        if not FORWARD:
            c0_pre = aux_desc.load([offs_am_c, offs_bn_c])
            c0 = 2 * c0 * tl.where(c0_pre > 0, c0_pre, 0)

        c_desc.store([offs_am_c, offs_bn_c], c0)

        if FORWARD:
            c0_post = tl.maximum(c0, 0)
            c0_post = c0_post * c0_post
            aux_desc.store([offs_am_c, offs_bn_c], c0_post)

        c1 = acc1.to(dtype)
        if not FORWARD:
            c1_pre = aux_desc.load([offs_am_c, offs_bn_c + BLOCK_SIZE_N // 2])
            c1 = 2 * c1 * tl.where(c1_pre > 0, c1_pre, 0)

        c_desc.store([offs_am_c, offs_bn_c + BLOCK_SIZE_N // 2], c1)

        if FORWARD:
            c1_post = tl.maximum(c1, 0)
            c1_post = c1_post * c1_post
            aux_desc.store([offs_am_c, offs_bn_c + BLOCK_SIZE_N // 2], c1_post)


def linear_relu_square(a, b, aux=None):
    M, K = a.shape
    N, K = b.shape
    dtype = a.dtype

    c = torch.empty((M, N), device=a.device, dtype=dtype)

    FORWARD = False
    if aux is None:
        FORWARD = True
        aux = torch.empty((M, N), device=a.device, dtype=dtype)

    NUM_SMS = torch.cuda.get_device_properties("cuda").multi_processor_count

    BLOCK_SIZE_M = 128
    BLOCK_SIZE_N = 256
    BLOCK_SIZE_K = 64
    num_stages = 4 if FORWARD else 3
    num_warps = 8

    a_desc = TensorDescriptor.from_tensor(a, [BLOCK_SIZE_M, BLOCK_SIZE_K])
    b_desc = TensorDescriptor.from_tensor(b, [BLOCK_SIZE_N, BLOCK_SIZE_K])
    c_desc = TensorDescriptor.from_tensor(c, [BLOCK_SIZE_M, BLOCK_SIZE_N // 2])
    aux_desc = TensorDescriptor.from_tensor(aux, [BLOCK_SIZE_M, BLOCK_SIZE_N // 2])

    def grid(META):
        return (min(
            NUM_SMS,
            triton.cdiv(M, BLOCK_SIZE_M) * triton.cdiv(N, BLOCK_SIZE_N),
        ), )

    linear_relu_square_kernel[grid](
        a_desc, b_desc, c_desc, aux_desc,
        M, N, K,
        BLOCK_SIZE_M=BLOCK_SIZE_M,
        BLOCK_SIZE_N=BLOCK_SIZE_N,
        BLOCK_SIZE_K=BLOCK_SIZE_K,
        GROUP_SIZE_M=1,
        NUM_SMS=NUM_SMS,
        FORWARD=FORWARD,
        num_stages=num_stages,
        num_warps=num_warps
    )

    if FORWARD:
        return c, aux
    else:
        return c

class FusedLinearReLUSquareFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, W1, W2):
        pre, post = linear_relu_square(x.view((-1, x.shape[-1])), W1)
        x3 = post @ W2
        ctx.save_for_backward(x, W1, W2, pre, post)
        return x3.view(x.shape)

    @staticmethod
    def backward(ctx, grad_output):
        x, W1, W2, pre, post = ctx.saved_tensors
        dW2 = post.T @ grad_output
        dpre = linear_relu_square(grad_output.view((-1, grad_output.shape[-1])), W2, aux=pre)
        dW1 = dpre.T @ x
        dx = dpre @ W1
        return dx.view(x.shape), dW1, dW2

# -----------------------------------------------------------------------------
# Fused Softcapped Cross Entropy


@triton.jit
def fused_softcapped_entropy_fwd_kernel(
    logits_ptr, losses_ptr, lse_ptr, targets_ptr, mtp_weights_ptr,
    stride_logits_n, stride_logits_v,
    n_rows, n_cols, n_predict,
    A, B, C,
    BLOCK_SIZE: tl.constexpr
):
    row_idx = tl.program_id(0).to(tl.int64)
    logits_row_ptr = logits_ptr + row_idx * stride_logits_n
    
    max_val = -float('inf')
    sum_exp = 0.0
    
    for off in range(0, n_cols, BLOCK_SIZE):
        cols = off + tl.arange(0, BLOCK_SIZE)
        mask = cols < n_cols
        val = tl.load(logits_row_ptr + cols, mask=mask, other=-float('inf')).to(tl.float32)
        z = A * tl.sigmoid((val + B) / C)
        z = tl.where(mask, z, -float('inf'))
        curr_max = tl.max(z, axis=0)
        new_max = tl.maximum(max_val, curr_max)
        sum_exp = sum_exp * tl.exp(max_val - new_max) + tl.sum(tl.exp(z - new_max), axis=0)
        max_val = new_max
    
    lse = max_val + tl.log(sum_exp)
    tl.store(lse_ptr + row_idx, lse)
    
    total_loss = 0.0
    for k in range(n_predict):
        target_idx = row_idx + k
        if target_idx < n_rows:
            weight = tl.load(mtp_weights_ptr + k)
            if weight > 0:
                target = tl.load(targets_ptr + target_idx).to(tl.int32)
                if target >= 0 and target < n_cols:
                    val_target = tl.load(logits_row_ptr + target).to(tl.float32)
                    z_target = A * tl.sigmoid((val_target + B) / C)
                    total_loss += weight * (lse - z_target)
    
    tl.store(losses_ptr + row_idx, total_loss)

@triton.jit
def fused_softcapped_entropy_bwd_kernel(
    grad_input_ptr, grad_output_ptr, lse_ptr, logits_ptr, targets_ptr, mtp_weights_ptr,
    stride_logits_n, stride_logits_v, stride_grad_n, stride_grad_v,
    n_rows, n_cols, n_predict,
    A, B, C,
    BLOCK_SIZE: tl.constexpr
):
    row_idx = tl.program_id(0).to(tl.int64)

    logits_row_ptr = logits_ptr + row_idx * stride_logits_n
    grad_row_ptr = grad_input_ptr + row_idx * stride_grad_n
    
    lse = tl.load(lse_ptr + row_idx)
    grad_loss = tl.load(grad_output_ptr + row_idx)
    
    S_w = 0.0
    for k in range(n_predict):
        if row_idx + k < n_rows:
            S_w += tl.load(mtp_weights_ptr + k)
            
    for off in range(0, n_cols, BLOCK_SIZE):
        cols = off + tl.arange(0, BLOCK_SIZE)
        mask = cols < n_cols
        val = tl.load(logits_row_ptr + cols, mask=mask, other=0.0).to(tl.float32)
        u = (val + B) / C
        sigmoid_u = tl.sigmoid(u)
        z = A * sigmoid_u
        p = tl.exp(z - lse)
        
        term1 = S_w * p
        term2 = tl.zeros([BLOCK_SIZE], dtype=tl.float32)
        for k in range(n_predict):
            if row_idx + k < n_rows:
                target = tl.load(targets_ptr + row_idx + k).to(tl.int32)
                weight = tl.load(mtp_weights_ptr + k)
                term2 += tl.where(cols == target, weight, 0.0)
        
        grad_z = grad_loss * (term1 - term2)
        dz_dx = (1.0 / C) * z * (1.0 - sigmoid_u)
        grad_x = grad_z * dz_dx
        tl.store(grad_row_ptr + cols, grad_x.to(tl.bfloat16), mask=mask)

class FusedSoftcappedCrossEntropy(torch.autograd.Function):
    @staticmethod
    def forward(ctx, logits, targets, mtp_weights, A=23.0, B=5.0, C=7.5):
        n_rows, n_cols = logits.shape
        if mtp_weights is None:
             mtp_weights = torch.tensor([1.0], device=logits.device, dtype=torch.float32)
        n_predict = mtp_weights.shape[0]

        losses = torch.empty(n_rows, dtype=torch.float32, device=logits.device)
        lse = torch.empty(n_rows, dtype=torch.float32, device=logits.device)
        
        logits = logits.contiguous()
        targets = targets.contiguous()
        mtp_weights = mtp_weights.contiguous()

        grid = (n_rows,)
        fused_softcapped_entropy_fwd_kernel[grid](
            logits, losses, lse, targets, mtp_weights,
            logits.stride(0), logits.stride(1),
            n_rows, n_cols, n_predict,
            A, B, C,
            BLOCK_SIZE=1024,
            num_warps=8,
            num_stages=4
        )
        
        ctx.save_for_backward(logits, targets, mtp_weights, lse)
        ctx.params = (A, B, C)
        return losses

    @staticmethod
    def backward(ctx, grad_output):
        logits, targets, mtp_weights, lse = ctx.saved_tensors
        A, B, C = ctx.params
        n_rows, n_cols = logits.shape
        n_predict = mtp_weights.shape[0]
        
        grad_input = torch.empty((n_rows, n_cols), dtype=torch.bfloat16, device=logits.device)
        grad_output = grad_output.contiguous()
        
        grid = (n_rows,)
        fused_softcapped_entropy_bwd_kernel[grid](
            grad_input, grad_output, lse, logits, targets, mtp_weights,
            logits.stride(0), logits.stride(1), grad_input.stride(0), grad_input.stride(1),
            n_rows, n_cols, n_predict,
            A, B, C,
            BLOCK_SIZE=1024,
            num_warps=8,
            num_stages=4
        )
        return grad_input, None, None, None, None, None
====================================================================================================
Running Python 3.10.12 (main, May 27 2025, 17:12:29) [GCC 11.4.0]
Running PyTorch 2.10.0.dev20251210+cu126 compiled for CUDA 12.6
Running Triton version 3.6.0
Mon Jan 26 02:40:14 2026       
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 570.148.08             Driver Version: 570.148.08     CUDA Version: 12.8     |
|-----------------------------------------+------------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  NVIDIA H100 80GB HBM3          On  |   00000000:61:00.0 Off |                    0 |
| N/A   35C    P0            121W /  700W |    1519MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   1  NVIDIA H100 80GB HBM3          On  |   00000000:62:00.0 Off |                    0 |
| N/A   38C    P0            121W /  700W |    1519MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   2  NVIDIA H100 80GB HBM3          On  |   00000000:63:00.0 Off |                    0 |
| N/A   41C    P0            121W /  700W |    1519MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   3  NVIDIA H100 80GB HBM3          On  |   00000000:64:00.0 Off |                    0 |
| N/A   35C    P0            119W /  700W |    1519MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   4  NVIDIA H100 80GB HBM3          On  |   00000000:6A:00.0 Off |                    0 |
| N/A   35C    P0            119W /  700W |    1519MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   5  NVIDIA H100 80GB HBM3          On  |   00000000:6B:00.0 Off |                    0 |
| N/A   42C    P0            132W /  700W |    1519MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   6  NVIDIA H100 80GB HBM3          On  |   00000000:6C:00.0 Off |                    0 |
| N/A   40C    P0            124W /  700W |    1519MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   7  NVIDIA H100 80GB HBM3          On  |   00000000:6D:00.0 Off |                    0 |
| N/A   36C    P0            118W /  700W |    1519MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
                                                                                         
+-----------------------------------------------------------------------------------------+
| Processes:                                                                              |
|  GPU   GI   CI              PID   Type   Process name                        GPU Memory |
|        ID   ID                                                               Usage      |
|=========================================================================================|
|    0   N/A  N/A          260848      C   /usr/bin/python3                       1510MiB |
|    1   N/A  N/A          260849      C   /usr/bin/python3                       1510MiB |
|    2   N/A  N/A          260850      C   /usr/bin/python3                       1510MiB |
|    3   N/A  N/A          260851      C   /usr/bin/python3                       1510MiB |
|    4   N/A  N/A          260852      C   /usr/bin/python3                       1510MiB |
|    5   N/A  N/A          260853      C   /usr/bin/python3                       1510MiB |
|    6   N/A  N/A          260854      C   /usr/bin/python3                       1510MiB |
|    7   N/A  N/A          260855      C   /usr/bin/python3                       1510MiB |
+-----------------------------------------------------------------------------------------+

====================================================================================================
Compiling model and warming up kernels (~7 minutes on first execution)
Sampling steps [0, 1, 2, 511, 512, 513, 1023, 1024, 1025, 1534, 1535, 1536] for warmup
Resetting Model
step:0/1575 val_loss:10.8318 train_time:0ms step_avg:0.03ms
step:1/1575 train_time:83ms step_avg:83.46ms
step:2/1575 train_time:104ms step_avg:51.93ms
step:3/1575 train_time:122ms step_avg:40.65ms
step:4/1575 train_time:162ms step_avg:40.39ms
step:5/1575 train_time:192ms step_avg:38.45ms
step:6/1575 train_time:292ms step_avg:48.68ms
step:7/1575 train_time:309ms step_avg:44.16ms
step:8/1575 train_time:333ms step_avg:41.65ms
step:9/1575 train_time:364ms step_avg:40.45ms
step:10/1575 train_time:402ms step_avg:40.21ms
step:11/1575 train_time:434ms step_avg:39.44ms
step:12/1575 train_time:472ms step_avg:39.37ms
step:13/1575 train_time:503ms step_avg:38.72ms
step:14/1575 train_time:542ms step_avg:38.71ms
step:15/1575 train_time:573ms step_avg:38.20ms
step:16/1575 train_time:612ms step_avg:38.22ms
step:17/1575 train_time:643ms step_avg:37.80ms
step:18/1575 train_time:681ms step_avg:37.85ms
step:19/1575 train_time:713ms step_avg:37.51ms
step:20/1575 train_time:751ms step_avg:37.57ms
step:21/1575 train_time:782ms step_avg:37.25ms
step:22/1575 train_time:821ms step_avg:37.32ms
step:23/1575 train_time:852ms step_avg:37.05ms
step:24/1575 train_time:891ms step_avg:37.12ms
step:25/1575 train_time:922ms step_avg:36.88ms
step:26/1575 train_time:960ms step_avg:36.94ms
step:27/1575 train_time:992ms step_avg:36.73ms
step:28/1575 train_time:1030ms step_avg:36.79ms
step:29/1575 train_time:1061ms step_avg:36.60ms
step:30/1575 train_time:1100ms step_avg:36.65ms
step:31/1575 train_time:1131ms step_avg:36.48ms
step:32/1575 train_time:1169ms step_avg:36.54ms
step:33/1575 train_time:1201ms step_avg:36.38ms
step:34/1575 train_time:1238ms step_avg:36.43ms
step:35/1575 train_time:1270ms step_avg:36.29ms
step:36/1575 train_time:1309ms step_avg:36.37ms
step:37/1575 train_time:1340ms step_avg:36.22ms
step:38/1575 train_time:1379ms step_avg:36.28ms
step:39/1575 train_time:1410ms step_avg:36.15ms
step:40/1575 train_time:1449ms step_avg:36.22ms
step:41/1575 train_time:1479ms step_avg:36.08ms
step:42/1575 train_time:1518ms step_avg:36.15ms
step:43/1575 train_time:1550ms step_avg:36.04ms
step:44/1575 train_time:1588ms step_avg:36.10ms
step:45/1575 train_time:1620ms step_avg:35.99ms
step:46/1575 train_time:1658ms step_avg:36.05ms
step:47/1575 train_time:1690ms step_avg:35.95ms
step:48/1575 train_time:1728ms step_avg:36.00ms
step:49/1575 train_time:1759ms step_avg:35.90ms
step:50/1575 train_time:1798ms step_avg:35.95ms
step:51/1575 train_time:1829ms step_avg:35.87ms
step:52/1575 train_time:1868ms step_avg:35.92ms
step:53/1575 train_time:1899ms step_avg:35.82ms
step:54/1575 train_time:1937ms step_avg:35.87ms
step:55/1575 train_time:1968ms step_avg:35.79ms
step:56/1575 train_time:2007ms step_avg:35.84ms
step:57/1575 train_time:2038ms step_avg:35.75ms
step:58/1575 train_time:2077ms step_avg:35.81ms
step:59/1575 train_time:2108ms step_avg:35.73ms
step:60/1575 train_time:2146ms step_avg:35.77ms
step:61/1575 train_time:2178ms step_avg:35.70ms
step:62/1575 train_time:2216ms step_avg:35.75ms
step:63/1575 train_time:2248ms step_avg:35.67ms
step:64/1575 train_time:2286ms step_avg:35.72ms
step:65/1575 train_time:2317ms step_avg:35.65ms
step:66/1575 train_time:2356ms step_avg:35.69ms
step:67/1575 train_time:2386ms step_avg:35.62ms
step:68/1575 train_time:2425ms step_avg:35.66ms
step:69/1575 train_time:2456ms step_avg:35.60ms
step:70/1575 train_time:2495ms step_avg:35.64ms
step:71/1575 train_time:2526ms step_avg:35.57ms
step:72/1575 train_time:2565ms step_avg:35.62ms
step:73/1575 train_time:2596ms step_avg:35.56ms
step:74/1575 train_time:2634ms step_avg:35.59ms
step:75/1575 train_time:2665ms step_avg:35.53ms
step:76/1575 train_time:2703ms step_avg:35.57ms
step:77/1575 train_time:2734ms step_avg:35.51ms
step:78/1575 train_time:2773ms step_avg:35.55ms
step:79/1575 train_time:2804ms step_avg:35.50ms
step:80/1575 train_time:2843ms step_avg:35.54ms
step:81/1575 train_time:2874ms step_avg:35.48ms
step:82/1575 train_time:2912ms step_avg:35.52ms
step:83/1575 train_time:2944ms step_avg:35.46ms
step:84/1575 train_time:2982ms step_avg:35.50ms
step:85/1575 train_time:3013ms step_avg:35.45ms
step:86/1575 train_time:3052ms step_avg:35.48ms
step:87/1575 train_time:3083ms step_avg:35.43ms
step:88/1575 train_time:3122ms step_avg:35.47ms
step:89/1575 train_time:3153ms step_avg:35.42ms
step:90/1575 train_time:3191ms step_avg:35.46ms
step:91/1575 train_time:3222ms step_avg:35.41ms
step:92/1575 train_time:3261ms step_avg:35.44ms
step:93/1575 train_time:3292ms step_avg:35.39ms
step:94/1575 train_time:3330ms step_avg:35.43ms
step:95/1575 train_time:3362ms step_avg:35.38ms
step:96/1575 train_time:3400ms step_avg:35.41ms
step:97/1575 train_time:3431ms step_avg:35.37ms
step:98/1575 train_time:3470ms step_avg:35.41ms
step:99/1575 train_time:3501ms step_avg:35.36ms
step:100/1575 train_time:3539ms step_avg:35.39ms
step:101/1575 train_time:3570ms step_avg:35.35ms
step:102/1575 train_time:3609ms step_avg:35.38ms
step:103/1575 train_time:3640ms step_avg:35.34ms
step:104/1575 train_time:3678ms step_avg:35.37ms
step:105/1575 train_time:3709ms step_avg:35.33ms
step:106/1575 train_time:3747ms step_avg:35.35ms
step:107/1575 train_time:3779ms step_avg:35.31ms
step:108/1575 train_time:3817ms step_avg:35.34ms
step:109/1575 train_time:3849ms step_avg:35.31ms
step:110/1575 train_time:3887ms step_avg:35.34ms
step:111/1575 train_time:3918ms step_avg:35.30ms
step:112/1575 train_time:3957ms step_avg:35.33ms
step:113/1575 train_time:3988ms step_avg:35.29ms
step:114/1575 train_time:4027ms step_avg:35.32ms
step:115/1575 train_time:4058ms step_avg:35.29ms
step:116/1575 train_time:4096ms step_avg:35.31ms
step:117/1575 train_time:4127ms step_avg:35.28ms
step:118/1575 train_time:4166ms step_avg:35.30ms
step:119/1575 train_time:4197ms step_avg:35.27ms
step:120/1575 train_time:4235ms step_avg:35.29ms
step:121/1575 train_time:4267ms step_avg:35.26ms
step:122/1575 train_time:4305ms step_avg:35.29ms
step:123/1575 train_time:4336ms step_avg:35.25ms
step:124/1575 train_time:4374ms step_avg:35.28ms
step:125/1575 train_time:4405ms step_avg:35.24ms
step:126/1575 train_time:4443ms step_avg:35.26ms
step:127/1575 train_time:4475ms step_avg:35.23ms
step:128/1575 train_time:4513ms step_avg:35.26ms
step:129/1575 train_time:4544ms step_avg:35.23ms
step:130/1575 train_time:4583ms step_avg:35.25ms
step:131/1575 train_time:4614ms step_avg:35.22ms
step:132/1575 train_time:4652ms step_avg:35.25ms
step:133/1575 train_time:4684ms step_avg:35.22ms
step:134/1575 train_time:4722ms step_avg:35.24ms
step:135/1575 train_time:4753ms step_avg:35.21ms
step:136/1575 train_time:4792ms step_avg:35.24ms
step:137/1575 train_time:4823ms step_avg:35.20ms
step:138/1575 train_time:4861ms step_avg:35.23ms
step:139/1575 train_time:4892ms step_avg:35.20ms
step:140/1575 train_time:4931ms step_avg:35.22ms
step:141/1575 train_time:4962ms step_avg:35.19ms
step:142/1575 train_time:5001ms step_avg:35.21ms
step:143/1575 train_time:5032ms step_avg:35.19ms
step:144/1575 train_time:5070ms step_avg:35.21ms
step:145/1575 train_time:5101ms step_avg:35.18ms
step:146/1575 train_time:5139ms step_avg:35.20ms
step:147/1575 train_time:5171ms step_avg:35.17ms
step:148/1575 train_time:5209ms step_avg:35.20ms
step:149/1575 train_time:5240ms step_avg:35.17ms
step:150/1575 train_time:5279ms step_avg:35.19ms
step:151/1575 train_time:5310ms step_avg:35.17ms
step:152/1575 train_time:5348ms step_avg:35.19ms
step:153/1575 train_time:5379ms step_avg:35.16ms
step:154/1575 train_time:5418ms step_avg:35.18ms
step:155/1575 train_time:5449ms step_avg:35.15ms
step:156/1575 train_time:5487ms step_avg:35.17ms
step:157/1575 train_time:5518ms step_avg:35.15ms
step:158/1575 train_time:5557ms step_avg:35.17ms
step:159/1575 train_time:5588ms step_avg:35.14ms
step:160/1575 train_time:5626ms step_avg:35.16ms
step:161/1575 train_time:5657ms step_avg:35.14ms
step:162/1575 train_time:5696ms step_avg:35.16ms
step:163/1575 train_time:5727ms step_avg:35.13ms
step:164/1575 train_time:5765ms step_avg:35.15ms
step:165/1575 train_time:5796ms step_avg:35.13ms
step:166/1575 train_time:5835ms step_avg:35.15ms
step:167/1575 train_time:5866ms step_avg:35.13ms
step:168/1575 train_time:5905ms step_avg:35.15ms
step:169/1575 train_time:5935ms step_avg:35.12ms
step:170/1575 train_time:5974ms step_avg:35.14ms
step:171/1575 train_time:6005ms step_avg:35.12ms
step:172/1575 train_time:6043ms step_avg:35.14ms
step:173/1575 train_time:6074ms step_avg:35.11ms
step:174/1575 train_time:6112ms step_avg:35.13ms
step:175/1575 train_time:6143ms step_avg:35.10ms
step:176/1575 train_time:6182ms step_avg:35.12ms
step:177/1575 train_time:6212ms step_avg:35.10ms
step:178/1575 train_time:6251ms step_avg:35.12ms
step:179/1575 train_time:6282ms step_avg:35.09ms
step:180/1575 train_time:6320ms step_avg:35.11ms
step:181/1575 train_time:6352ms step_avg:35.09ms
step:182/1575 train_time:6390ms step_avg:35.11ms
step:183/1575 train_time:6422ms step_avg:35.09ms
step:184/1575 train_time:6460ms step_avg:35.11ms
step:185/1575 train_time:6491ms step_avg:35.09ms
step:186/1575 train_time:6530ms step_avg:35.11ms
step:187/1575 train_time:6561ms step_avg:35.09ms
step:188/1575 train_time:6599ms step_avg:35.10ms
step:189/1575 train_time:6630ms step_avg:35.08ms
step:190/1575 train_time:6670ms step_avg:35.10ms
step:191/1575 train_time:6700ms step_avg:35.08ms
step:192/1575 train_time:6738ms step_avg:35.09ms
step:193/1575 train_time:6769ms step_avg:35.07ms
step:194/1575 train_time:6807ms step_avg:35.09ms
step:195/1575 train_time:6838ms step_avg:35.07ms
step:196/1575 train_time:6877ms step_avg:35.09ms
step:197/1575 train_time:6908ms step_avg:35.06ms
step:198/1575 train_time:6946ms step_avg:35.08ms
step:199/1575 train_time:6977ms step_avg:35.06ms
step:200/1575 train_time:7016ms step_avg:35.08ms
step:201/1575 train_time:7047ms step_avg:35.06ms
step:202/1575 train_time:7085ms step_avg:35.07ms
step:203/1575 train_time:7116ms step_avg:35.06ms
step:204/1575 train_time:7155ms step_avg:35.07ms
step:205/1575 train_time:7186ms step_avg:35.05ms
step:206/1575 train_time:7225ms step_avg:35.07ms
step:207/1575 train_time:7255ms step_avg:35.05ms
step:208/1575 train_time:7296ms step_avg:35.08ms
step:209/1575 train_time:7325ms step_avg:35.05ms
step:210/1575 train_time:7363ms step_avg:35.06ms
step:211/1575 train_time:7394ms step_avg:35.04ms
step:212/1575 train_time:7432ms step_avg:35.06ms
step:213/1575 train_time:7463ms step_avg:35.04ms
step:214/1575 train_time:7502ms step_avg:35.06ms
step:215/1575 train_time:7533ms step_avg:35.04ms
step:216/1575 train_time:7572ms step_avg:35.06ms
step:217/1575 train_time:7602ms step_avg:35.03ms
step:218/1575 train_time:7641ms step_avg:35.05ms
step:219/1575 train_time:7672ms step_avg:35.03ms
step:220/1575 train_time:7710ms step_avg:35.04ms
step:221/1575 train_time:7741ms step_avg:35.03ms
step:222/1575 train_time:7779ms step_avg:35.04ms
step:223/1575 train_time:7811ms step_avg:35.02ms
step:224/1575 train_time:7849ms step_avg:35.04ms
step:225/1575 train_time:7880ms step_avg:35.02ms
step:226/1575 train_time:7918ms step_avg:35.04ms
step:227/1575 train_time:7950ms step_avg:35.02ms
step:228/1575 train_time:7988ms step_avg:35.04ms
step:229/1575 train_time:8019ms step_avg:35.02ms
step:230/1575 train_time:8058ms step_avg:35.03ms
step:231/1575 train_time:8089ms step_avg:35.02ms
step:232/1575 train_time:8127ms step_avg:35.03ms
step:233/1575 train_time:8158ms step_avg:35.01ms
step:234/1575 train_time:8196ms step_avg:35.03ms
step:235/1575 train_time:8227ms step_avg:35.01ms
step:236/1575 train_time:8266ms step_avg:35.02ms
step:237/1575 train_time:8297ms step_avg:35.01ms
step:238/1575 train_time:8335ms step_avg:35.02ms
step:239/1575 train_time:8366ms step_avg:35.00ms
step:240/1575 train_time:8404ms step_avg:35.02ms
step:241/1575 train_time:8436ms step_avg:35.00ms
step:242/1575 train_time:8474ms step_avg:35.02ms
step:243/1575 train_time:8505ms step_avg:35.00ms
step:244/1575 train_time:8543ms step_avg:35.01ms
step:245/1575 train_time:8575ms step_avg:35.00ms
step:246/1575 train_time:8613ms step_avg:35.01ms
step:247/1575 train_time:8644ms step_avg:35.00ms
step:248/1575 train_time:8682ms step_avg:35.01ms
step:249/1575 train_time:8714ms step_avg:34.99ms
step:250/1575 train_time:8752ms step_avg:35.01ms
step:250/1575 val_loss:4.5735 train_time:8801ms step_avg:35.20ms
step:251/1575 train_time:8819ms step_avg:35.14ms
step:252/1575 train_time:8838ms step_avg:35.07ms
step:253/1575 train_time:8858ms step_avg:35.01ms
step:254/1575 train_time:8897ms step_avg:35.03ms
step:255/1575 train_time:8930ms step_avg:35.02ms
step:256/1575 train_time:8970ms step_avg:35.04ms
step:257/1575 train_time:9002ms step_avg:35.03ms
step:258/1575 train_time:9040ms step_avg:35.04ms
step:259/1575 train_time:9071ms step_avg:35.02ms
step:260/1575 train_time:9110ms step_avg:35.04ms
step:261/1575 train_time:9141ms step_avg:35.02ms
step:262/1575 train_time:9179ms step_avg:35.03ms
step:263/1575 train_time:9210ms step_avg:35.02ms
step:264/1575 train_time:9249ms step_avg:35.03ms
step:265/1575 train_time:9280ms step_avg:35.02ms
step:266/1575 train_time:9318ms step_avg:35.03ms
step:267/1575 train_time:9349ms step_avg:35.02ms
step:268/1575 train_time:9387ms step_avg:35.03ms
step:269/1575 train_time:9419ms step_avg:35.01ms
step:270/1575 train_time:9457ms step_avg:35.03ms
step:271/1575 train_time:9488ms step_avg:35.01ms
step:272/1575 train_time:9527ms step_avg:35.03ms
step:273/1575 train_time:9557ms step_avg:35.01ms
step:274/1575 train_time:9596ms step_avg:35.02ms
step:275/1575 train_time:9627ms step_avg:35.01ms
step:276/1575 train_time:9665ms step_avg:35.02ms
step:277/1575 train_time:9696ms step_avg:35.00ms
step:278/1575 train_time:9734ms step_avg:35.02ms
step:279/1575 train_time:9765ms step_avg:35.00ms
step:280/1575 train_time:9803ms step_avg:35.01ms
step:281/1575 train_time:9834ms step_avg:35.00ms
step:282/1575 train_time:9873ms step_avg:35.01ms
step:283/1575 train_time:9904ms step_avg:35.00ms
step:284/1575 train_time:9943ms step_avg:35.01ms
step:285/1575 train_time:9973ms step_avg:34.99ms
step:286/1575 train_time:10012ms step_avg:35.01ms
step:287/1575 train_time:10043ms step_avg:34.99ms
step:288/1575 train_time:10082ms step_avg:35.01ms
step:289/1575 train_time:10113ms step_avg:34.99ms
step:290/1575 train_time:10152ms step_avg:35.01ms
step:291/1575 train_time:10183ms step_avg:34.99ms
step:292/1575 train_time:10221ms step_avg:35.00ms
step:293/1575 train_time:10252ms step_avg:34.99ms
step:294/1575 train_time:10290ms step_avg:35.00ms
step:295/1575 train_time:10322ms step_avg:34.99ms
step:296/1575 train_time:10360ms step_avg:35.00ms
step:297/1575 train_time:10391ms step_avg:34.99ms
step:298/1575 train_time:10429ms step_avg:35.00ms
step:299/1575 train_time:10461ms step_avg:34.99ms
step:300/1575 train_time:10499ms step_avg:35.00ms
step:301/1575 train_time:10530ms step_avg:34.98ms
step:302/1575 train_time:10568ms step_avg:34.99ms
step:303/1575 train_time:10600ms step_avg:34.98ms
step:304/1575 train_time:10638ms step_avg:34.99ms
step:305/1575 train_time:10669ms step_avg:34.98ms
step:306/1575 train_time:10707ms step_avg:34.99ms
step:307/1575 train_time:10738ms step_avg:34.98ms
step:308/1575 train_time:10777ms step_avg:34.99ms
step:309/1575 train_time:10808ms step_avg:34.98ms
step:310/1575 train_time:10846ms step_avg:34.99ms
step:311/1575 train_time:10877ms step_avg:34.97ms
step:312/1575 train_time:10916ms step_avg:34.99ms
step:313/1575 train_time:10947ms step_avg:34.97ms
step:314/1575 train_time:10985ms step_avg:34.99ms
step:315/1575 train_time:11017ms step_avg:34.97ms
step:316/1575 train_time:11055ms step_avg:34.98ms
step:317/1575 train_time:11086ms step_avg:34.97ms
step:318/1575 train_time:11124ms step_avg:34.98ms
step:319/1575 train_time:11155ms step_avg:34.97ms
step:320/1575 train_time:11194ms step_avg:34.98ms
step:321/1575 train_time:11225ms step_avg:34.97ms
step:322/1575 train_time:11264ms step_avg:34.98ms
step:323/1575 train_time:11294ms step_avg:34.97ms
step:324/1575 train_time:11333ms step_avg:34.98ms
step:325/1575 train_time:11364ms step_avg:34.97ms
step:326/1575 train_time:11402ms step_avg:34.98ms
step:327/1575 train_time:11433ms step_avg:34.96ms
step:328/1575 train_time:11471ms step_avg:34.97ms
step:329/1575 train_time:11502ms step_avg:34.96ms
step:330/1575 train_time:11541ms step_avg:34.97ms
step:331/1575 train_time:11572ms step_avg:34.96ms
step:332/1575 train_time:11610ms step_avg:34.97ms
step:333/1575 train_time:11642ms step_avg:34.96ms
step:334/1575 train_time:11680ms step_avg:34.97ms
step:335/1575 train_time:11711ms step_avg:34.96ms
step:336/1575 train_time:11750ms step_avg:34.97ms
step:337/1575 train_time:11781ms step_avg:34.96ms
step:338/1575 train_time:11819ms step_avg:34.97ms
step:339/1575 train_time:11850ms step_avg:34.96ms
step:340/1575 train_time:11889ms step_avg:34.97ms
step:341/1575 train_time:11920ms step_avg:34.96ms
step:342/1575 train_time:11959ms step_avg:34.97ms
step:343/1575 train_time:11990ms step_avg:34.96ms
step:344/1575 train_time:12029ms step_avg:34.97ms
step:345/1575 train_time:12060ms step_avg:34.96ms
step:346/1575 train_time:12098ms step_avg:34.97ms
step:347/1575 train_time:12129ms step_avg:34.96ms
step:348/1575 train_time:12168ms step_avg:34.96ms
step:349/1575 train_time:12199ms step_avg:34.95ms
step:350/1575 train_time:12237ms step_avg:34.96ms
step:351/1575 train_time:12268ms step_avg:34.95ms
step:352/1575 train_time:12306ms step_avg:34.96ms
step:353/1575 train_time:12337ms step_avg:34.95ms
step:354/1575 train_time:12376ms step_avg:34.96ms
step:355/1575 train_time:12407ms step_avg:34.95ms
step:356/1575 train_time:12445ms step_avg:34.96ms
step:357/1575 train_time:12476ms step_avg:34.95ms
step:358/1575 train_time:12515ms step_avg:34.96ms
step:359/1575 train_time:12546ms step_avg:34.95ms
step:360/1575 train_time:12584ms step_avg:34.95ms
step:361/1575 train_time:12614ms step_avg:34.94ms
step:362/1575 train_time:12653ms step_avg:34.95ms
step:363/1575 train_time:12684ms step_avg:34.94ms
step:364/1575 train_time:12722ms step_avg:34.95ms
step:365/1575 train_time:12753ms step_avg:34.94ms
step:366/1575 train_time:12791ms step_avg:34.95ms
step:367/1575 train_time:12822ms step_avg:34.94ms
step:368/1575 train_time:12860ms step_avg:34.95ms
step:369/1575 train_time:12891ms step_avg:34.94ms
step:370/1575 train_time:12930ms step_avg:34.94ms
step:371/1575 train_time:12961ms step_avg:34.94ms
step:372/1575 train_time:12999ms step_avg:34.94ms
step:373/1575 train_time:13031ms step_avg:34.93ms
step:374/1575 train_time:13069ms step_avg:34.94ms
step:375/1575 train_time:13101ms step_avg:34.93ms
step:376/1575 train_time:13140ms step_avg:34.95ms
step:377/1575 train_time:13170ms step_avg:34.93ms
step:378/1575 train_time:13209ms step_avg:34.94ms
step:379/1575 train_time:13240ms step_avg:34.93ms
step:380/1575 train_time:13278ms step_avg:34.94ms
step:381/1575 train_time:13309ms step_avg:34.93ms
step:382/1575 train_time:13348ms step_avg:34.94ms
step:383/1575 train_time:13379ms step_avg:34.93ms
step:384/1575 train_time:13418ms step_avg:34.94ms
step:385/1575 train_time:13449ms step_avg:34.93ms
step:386/1575 train_time:13487ms step_avg:34.94ms
step:387/1575 train_time:13518ms step_avg:34.93ms
step:388/1575 train_time:13557ms step_avg:34.94ms
step:389/1575 train_time:13588ms step_avg:34.93ms
step:390/1575 train_time:13626ms step_avg:34.94ms
step:391/1575 train_time:13658ms step_avg:34.93ms
step:392/1575 train_time:13696ms step_avg:34.94ms
step:393/1575 train_time:13727ms step_avg:34.93ms
step:394/1575 train_time:13766ms step_avg:34.94ms
step:395/1575 train_time:13797ms step_avg:34.93ms
step:396/1575 train_time:13835ms step_avg:34.94ms
step:397/1575 train_time:13866ms step_avg:34.93ms
step:398/1575 train_time:13905ms step_avg:34.94ms
step:399/1575 train_time:13935ms step_avg:34.93ms
step:400/1575 train_time:13974ms step_avg:34.94ms
step:401/1575 train_time:14005ms step_avg:34.93ms
step:402/1575 train_time:14043ms step_avg:34.93ms
step:403/1575 train_time:14074ms step_avg:34.92ms
step:404/1575 train_time:14112ms step_avg:34.93ms
step:405/1575 train_time:14144ms step_avg:34.92ms
step:406/1575 train_time:14182ms step_avg:34.93ms
step:407/1575 train_time:14213ms step_avg:34.92ms
step:408/1575 train_time:14251ms step_avg:34.93ms
step:409/1575 train_time:14282ms step_avg:34.92ms
step:410/1575 train_time:14321ms step_avg:34.93ms
step:411/1575 train_time:14352ms step_avg:34.92ms
step:412/1575 train_time:14390ms step_avg:34.93ms
step:413/1575 train_time:14421ms step_avg:34.92ms
step:414/1575 train_time:14461ms step_avg:34.93ms
step:415/1575 train_time:14492ms step_avg:34.92ms
step:416/1575 train_time:14530ms step_avg:34.93ms
step:417/1575 train_time:14562ms step_avg:34.92ms
step:418/1575 train_time:14600ms step_avg:34.93ms
step:419/1575 train_time:14631ms step_avg:34.92ms
step:420/1575 train_time:14670ms step_avg:34.93ms
step:421/1575 train_time:14700ms step_avg:34.92ms
step:422/1575 train_time:14740ms step_avg:34.93ms
step:423/1575 train_time:14770ms step_avg:34.92ms
step:424/1575 train_time:14808ms step_avg:34.92ms
step:425/1575 train_time:14839ms step_avg:34.91ms
step:426/1575 train_time:14878ms step_avg:34.92ms
step:427/1575 train_time:14908ms step_avg:34.91ms
step:428/1575 train_time:14947ms step_avg:34.92ms
step:429/1575 train_time:14978ms step_avg:34.91ms
step:430/1575 train_time:15016ms step_avg:34.92ms
step:431/1575 train_time:15048ms step_avg:34.91ms
step:432/1575 train_time:15086ms step_avg:34.92ms
step:433/1575 train_time:15117ms step_avg:34.91ms
step:434/1575 train_time:15155ms step_avg:34.92ms
step:435/1575 train_time:15186ms step_avg:34.91ms
step:436/1575 train_time:15225ms step_avg:34.92ms
step:437/1575 train_time:15256ms step_avg:34.91ms
step:438/1575 train_time:15294ms step_avg:34.92ms
step:439/1575 train_time:15326ms step_avg:34.91ms
step:440/1575 train_time:15365ms step_avg:34.92ms
step:441/1575 train_time:15395ms step_avg:34.91ms
step:442/1575 train_time:15434ms step_avg:34.92ms
step:443/1575 train_time:15465ms step_avg:34.91ms
step:444/1575 train_time:15503ms step_avg:34.92ms
step:445/1575 train_time:15534ms step_avg:34.91ms
step:446/1575 train_time:15573ms step_avg:34.92ms
step:447/1575 train_time:15603ms step_avg:34.91ms
step:448/1575 train_time:15642ms step_avg:34.91ms
step:449/1575 train_time:15673ms step_avg:34.91ms
step:450/1575 train_time:15711ms step_avg:34.91ms
step:451/1575 train_time:15742ms step_avg:34.91ms
step:452/1575 train_time:15781ms step_avg:34.91ms
step:453/1575 train_time:15811ms step_avg:34.90ms
step:454/1575 train_time:15850ms step_avg:34.91ms
step:455/1575 train_time:15881ms step_avg:34.90ms
step:456/1575 train_time:15919ms step_avg:34.91ms
step:457/1575 train_time:15950ms step_avg:34.90ms
step:458/1575 train_time:15989ms step_avg:34.91ms
step:459/1575 train_time:16020ms step_avg:34.90ms
step:460/1575 train_time:16058ms step_avg:34.91ms
step:461/1575 train_time:16089ms step_avg:34.90ms
step:462/1575 train_time:16128ms step_avg:34.91ms
step:463/1575 train_time:16159ms step_avg:34.90ms
step:464/1575 train_time:16197ms step_avg:34.91ms
step:465/1575 train_time:16228ms step_avg:34.90ms
step:466/1575 train_time:16267ms step_avg:34.91ms
step:467/1575 train_time:16298ms step_avg:34.90ms
step:468/1575 train_time:16336ms step_avg:34.91ms
step:469/1575 train_time:16368ms step_avg:34.90ms
step:470/1575 train_time:16407ms step_avg:34.91ms
step:471/1575 train_time:16438ms step_avg:34.90ms
step:472/1575 train_time:16477ms step_avg:34.91ms
step:473/1575 train_time:16508ms step_avg:34.90ms
step:474/1575 train_time:16546ms step_avg:34.91ms
step:475/1575 train_time:16578ms step_avg:34.90ms
step:476/1575 train_time:16616ms step_avg:34.91ms
step:477/1575 train_time:16647ms step_avg:34.90ms
step:478/1575 train_time:16685ms step_avg:34.91ms
step:479/1575 train_time:16716ms step_avg:34.90ms
step:480/1575 train_time:16755ms step_avg:34.91ms
step:481/1575 train_time:16785ms step_avg:34.90ms
step:482/1575 train_time:16824ms step_avg:34.90ms
step:483/1575 train_time:16855ms step_avg:34.90ms
step:484/1575 train_time:16893ms step_avg:34.90ms
step:485/1575 train_time:16924ms step_avg:34.89ms
step:486/1575 train_time:16962ms step_avg:34.90ms
step:487/1575 train_time:16993ms step_avg:34.89ms
step:488/1575 train_time:17032ms step_avg:34.90ms
step:489/1575 train_time:17063ms step_avg:34.89ms
step:490/1575 train_time:17101ms step_avg:34.90ms
step:491/1575 train_time:17132ms step_avg:34.89ms
step:492/1575 train_time:17171ms step_avg:34.90ms
step:493/1575 train_time:17202ms step_avg:34.89ms
step:494/1575 train_time:17240ms step_avg:34.90ms
step:495/1575 train_time:17271ms step_avg:34.89ms
step:496/1575 train_time:17309ms step_avg:34.90ms
step:497/1575 train_time:17340ms step_avg:34.89ms
step:498/1575 train_time:17379ms step_avg:34.90ms
step:499/1575 train_time:17410ms step_avg:34.89ms
step:500/1575 train_time:17448ms step_avg:34.90ms
step:500/1575 val_loss:4.2405 train_time:17497ms step_avg:34.99ms
step:501/1575 train_time:17515ms step_avg:34.96ms
step:502/1575 train_time:17534ms step_avg:34.93ms
step:503/1575 train_time:17553ms step_avg:34.90ms
step:504/1575 train_time:17592ms step_avg:34.91ms
step:505/1575 train_time:17625ms step_avg:34.90ms
step:506/1575 train_time:17665ms step_avg:34.91ms
step:507/1575 train_time:17696ms step_avg:34.90ms
step:508/1575 train_time:17735ms step_avg:34.91ms
step:509/1575 train_time:17766ms step_avg:34.90ms
step:510/1575 train_time:17805ms step_avg:34.91ms
step:511/1575 train_time:17836ms step_avg:34.90ms
step:512/1575 train_time:17874ms step_avg:34.91ms
step:513/1575 train_time:17945ms step_avg:34.98ms
step:514/1575 train_time:18001ms step_avg:35.02ms
step:515/1575 train_time:18064ms step_avg:35.08ms
step:516/1575 train_time:18123ms step_avg:35.12ms
step:517/1575 train_time:18185ms step_avg:35.17ms
step:518/1575 train_time:18244ms step_avg:35.22ms
step:519/1575 train_time:18306ms step_avg:35.27ms
step:520/1575 train_time:18365ms step_avg:35.32ms
step:521/1575 train_time:18428ms step_avg:35.37ms
step:522/1575 train_time:18488ms step_avg:35.42ms
step:523/1575 train_time:18553ms step_avg:35.47ms
step:524/1575 train_time:18614ms step_avg:35.52ms
step:525/1575 train_time:18678ms step_avg:35.58ms
step:526/1575 train_time:18739ms step_avg:35.62ms
step:527/1575 train_time:18804ms step_avg:35.68ms
step:528/1575 train_time:18864ms step_avg:35.73ms
step:529/1575 train_time:18928ms step_avg:35.78ms
step:530/1575 train_time:18987ms step_avg:35.82ms
step:531/1575 train_time:19051ms step_avg:35.88ms
step:532/1575 train_time:19110ms step_avg:35.92ms
step:533/1575 train_time:19172ms step_avg:35.97ms
step:534/1575 train_time:19233ms step_avg:36.02ms
step:535/1575 train_time:19295ms step_avg:36.07ms
step:536/1575 train_time:19353ms step_avg:36.11ms
step:537/1575 train_time:19416ms step_avg:36.16ms
step:538/1575 train_time:19475ms step_avg:36.20ms
step:539/1575 train_time:19538ms step_avg:36.25ms
step:540/1575 train_time:19598ms step_avg:36.29ms
step:541/1575 train_time:19661ms step_avg:36.34ms
step:542/1575 train_time:19721ms step_avg:36.39ms
step:543/1575 train_time:19786ms step_avg:36.44ms
step:544/1575 train_time:19845ms step_avg:36.48ms
step:545/1575 train_time:19910ms step_avg:36.53ms
step:546/1575 train_time:19969ms step_avg:36.57ms
step:547/1575 train_time:20032ms step_avg:36.62ms
step:548/1575 train_time:20091ms step_avg:36.66ms
step:549/1575 train_time:20154ms step_avg:36.71ms
step:550/1575 train_time:20213ms step_avg:36.75ms
step:551/1575 train_time:20276ms step_avg:36.80ms
step:552/1575 train_time:20335ms step_avg:36.84ms
step:553/1575 train_time:20398ms step_avg:36.89ms
step:554/1575 train_time:20458ms step_avg:36.93ms
step:555/1575 train_time:20521ms step_avg:36.98ms
step:556/1575 train_time:20583ms step_avg:37.02ms
step:557/1575 train_time:20644ms step_avg:37.06ms
step:558/1575 train_time:20704ms step_avg:37.10ms
step:559/1575 train_time:20768ms step_avg:37.15ms
step:560/1575 train_time:20828ms step_avg:37.19ms
step:561/1575 train_time:20891ms step_avg:37.24ms
step:562/1575 train_time:20950ms step_avg:37.28ms
step:563/1575 train_time:21014ms step_avg:37.32ms
step:564/1575 train_time:21073ms step_avg:37.36ms
step:565/1575 train_time:21136ms step_avg:37.41ms
step:566/1575 train_time:21195ms step_avg:37.45ms
step:567/1575 train_time:21258ms step_avg:37.49ms
step:568/1575 train_time:21317ms step_avg:37.53ms
step:569/1575 train_time:21385ms step_avg:37.58ms
step:570/1575 train_time:21442ms step_avg:37.62ms
step:571/1575 train_time:21504ms step_avg:37.66ms
step:572/1575 train_time:21564ms step_avg:37.70ms
step:573/1575 train_time:21626ms step_avg:37.74ms
step:574/1575 train_time:21685ms step_avg:37.78ms
step:575/1575 train_time:21751ms step_avg:37.83ms
step:576/1575 train_time:21809ms step_avg:37.86ms
step:577/1575 train_time:21872ms step_avg:37.91ms
step:578/1575 train_time:21932ms step_avg:37.94ms
step:579/1575 train_time:21995ms step_avg:37.99ms
step:580/1575 train_time:22054ms step_avg:38.02ms
step:581/1575 train_time:22118ms step_avg:38.07ms
step:582/1575 train_time:22176ms step_avg:38.10ms
step:583/1575 train_time:22240ms step_avg:38.15ms
step:584/1575 train_time:22299ms step_avg:38.18ms
step:585/1575 train_time:22363ms step_avg:38.23ms
step:586/1575 train_time:22422ms step_avg:38.26ms
step:587/1575 train_time:22484ms step_avg:38.30ms
step:588/1575 train_time:22544ms step_avg:38.34ms
step:589/1575 train_time:22607ms step_avg:38.38ms
step:590/1575 train_time:22666ms step_avg:38.42ms
step:591/1575 train_time:22730ms step_avg:38.46ms
step:592/1575 train_time:22790ms step_avg:38.50ms
step:593/1575 train_time:22853ms step_avg:38.54ms
step:594/1575 train_time:22912ms step_avg:38.57ms
step:595/1575 train_time:22978ms step_avg:38.62ms
step:596/1575 train_time:23035ms step_avg:38.65ms
step:597/1575 train_time:23097ms step_avg:38.69ms
step:598/1575 train_time:23157ms step_avg:38.72ms
step:599/1575 train_time:23220ms step_avg:38.77ms
step:600/1575 train_time:23279ms step_avg:38.80ms
step:601/1575 train_time:23343ms step_avg:38.84ms
step:602/1575 train_time:23402ms step_avg:38.87ms
step:603/1575 train_time:23466ms step_avg:38.92ms
step:604/1575 train_time:23525ms step_avg:38.95ms
step:605/1575 train_time:23588ms step_avg:38.99ms
step:606/1575 train_time:23647ms step_avg:39.02ms
step:607/1575 train_time:23711ms step_avg:39.06ms
step:608/1575 train_time:23770ms step_avg:39.10ms
step:609/1575 train_time:23833ms step_avg:39.13ms
step:610/1575 train_time:23892ms step_avg:39.17ms
step:611/1575 train_time:23956ms step_avg:39.21ms
step:612/1575 train_time:24015ms step_avg:39.24ms
step:613/1575 train_time:24079ms step_avg:39.28ms
step:614/1575 train_time:24138ms step_avg:39.31ms
step:615/1575 train_time:24202ms step_avg:39.35ms
step:616/1575 train_time:24261ms step_avg:39.38ms
step:617/1575 train_time:24324ms step_avg:39.42ms
step:618/1575 train_time:24383ms step_avg:39.45ms
step:619/1575 train_time:24446ms step_avg:39.49ms
step:620/1575 train_time:24505ms step_avg:39.52ms
step:621/1575 train_time:24569ms step_avg:39.56ms
step:622/1575 train_time:24627ms step_avg:39.59ms
step:623/1575 train_time:24691ms step_avg:39.63ms
step:624/1575 train_time:24750ms step_avg:39.66ms
step:625/1575 train_time:24814ms step_avg:39.70ms
step:626/1575 train_time:24874ms step_avg:39.73ms
step:627/1575 train_time:24937ms step_avg:39.77ms
step:628/1575 train_time:24996ms step_avg:39.80ms
step:629/1575 train_time:25060ms step_avg:39.84ms
step:630/1575 train_time:25118ms step_avg:39.87ms
step:631/1575 train_time:25182ms step_avg:39.91ms
step:632/1575 train_time:25241ms step_avg:39.94ms
step:633/1575 train_time:25304ms step_avg:39.97ms
step:634/1575 train_time:25367ms step_avg:40.01ms
step:635/1575 train_time:25429ms step_avg:40.05ms
step:636/1575 train_time:25487ms step_avg:40.07ms
step:637/1575 train_time:25550ms step_avg:40.11ms
step:638/1575 train_time:25609ms step_avg:40.14ms
step:639/1575 train_time:25672ms step_avg:40.17ms
step:640/1575 train_time:25731ms step_avg:40.20ms
step:641/1575 train_time:25795ms step_avg:40.24ms
step:642/1575 train_time:25854ms step_avg:40.27ms
step:643/1575 train_time:25918ms step_avg:40.31ms
step:644/1575 train_time:25978ms step_avg:40.34ms
step:645/1575 train_time:26041ms step_avg:40.37ms
step:646/1575 train_time:26100ms step_avg:40.40ms
step:647/1575 train_time:26163ms step_avg:40.44ms
step:648/1575 train_time:26222ms step_avg:40.47ms
step:649/1575 train_time:26285ms step_avg:40.50ms
step:650/1575 train_time:26344ms step_avg:40.53ms
step:651/1575 train_time:26408ms step_avg:40.56ms
step:652/1575 train_time:26467ms step_avg:40.59ms
step:653/1575 train_time:26531ms step_avg:40.63ms
step:654/1575 train_time:26590ms step_avg:40.66ms
step:655/1575 train_time:26656ms step_avg:40.70ms
step:656/1575 train_time:26714ms step_avg:40.72ms
step:657/1575 train_time:26777ms step_avg:40.76ms
step:658/1575 train_time:26837ms step_avg:40.79ms
step:659/1575 train_time:26900ms step_avg:40.82ms
step:660/1575 train_time:26959ms step_avg:40.85ms
step:661/1575 train_time:27022ms step_avg:40.88ms
step:662/1575 train_time:27081ms step_avg:40.91ms
step:663/1575 train_time:27145ms step_avg:40.94ms
step:664/1575 train_time:27204ms step_avg:40.97ms
step:665/1575 train_time:27268ms step_avg:41.00ms
step:666/1575 train_time:27327ms step_avg:41.03ms
step:667/1575 train_time:27391ms step_avg:41.07ms
step:668/1575 train_time:27449ms step_avg:41.09ms
step:669/1575 train_time:27513ms step_avg:41.13ms
step:670/1575 train_time:27573ms step_avg:41.15ms
step:671/1575 train_time:27636ms step_avg:41.19ms
step:672/1575 train_time:27695ms step_avg:41.21ms
step:673/1575 train_time:27758ms step_avg:41.25ms
step:674/1575 train_time:27818ms step_avg:41.27ms
step:675/1575 train_time:27881ms step_avg:41.31ms
step:676/1575 train_time:27941ms step_avg:41.33ms
step:677/1575 train_time:28004ms step_avg:41.36ms
step:678/1575 train_time:28063ms step_avg:41.39ms
step:679/1575 train_time:28126ms step_avg:41.42ms
step:680/1575 train_time:28185ms step_avg:41.45ms
step:681/1575 train_time:28249ms step_avg:41.48ms
step:682/1575 train_time:28308ms step_avg:41.51ms
step:683/1575 train_time:28371ms step_avg:41.54ms
step:684/1575 train_time:28430ms step_avg:41.56ms
step:685/1575 train_time:28494ms step_avg:41.60ms
step:686/1575 train_time:28553ms step_avg:41.62ms
step:687/1575 train_time:28616ms step_avg:41.65ms
step:688/1575 train_time:28676ms step_avg:41.68ms
step:689/1575 train_time:28738ms step_avg:41.71ms
step:690/1575 train_time:28797ms step_avg:41.74ms
step:691/1575 train_time:28861ms step_avg:41.77ms
step:692/1575 train_time:28920ms step_avg:41.79ms
step:693/1575 train_time:28984ms step_avg:41.82ms
step:694/1575 train_time:29043ms step_avg:41.85ms
step:695/1575 train_time:29107ms step_avg:41.88ms
step:696/1575 train_time:29167ms step_avg:41.91ms
step:697/1575 train_time:29230ms step_avg:41.94ms
step:698/1575 train_time:29290ms step_avg:41.96ms
step:699/1575 train_time:29353ms step_avg:41.99ms
step:700/1575 train_time:29414ms step_avg:42.02ms
step:701/1575 train_time:29476ms step_avg:42.05ms
step:702/1575 train_time:29535ms step_avg:42.07ms
step:703/1575 train_time:29599ms step_avg:42.10ms
step:704/1575 train_time:29658ms step_avg:42.13ms
step:705/1575 train_time:29721ms step_avg:42.16ms
step:706/1575 train_time:29780ms step_avg:42.18ms
step:707/1575 train_time:29843ms step_avg:42.21ms
step:708/1575 train_time:29902ms step_avg:42.23ms
step:709/1575 train_time:29966ms step_avg:42.26ms
step:710/1575 train_time:30025ms step_avg:42.29ms
step:711/1575 train_time:30088ms step_avg:42.32ms
step:712/1575 train_time:30147ms step_avg:42.34ms
step:713/1575 train_time:30211ms step_avg:42.37ms
step:714/1575 train_time:30270ms step_avg:42.39ms
step:715/1575 train_time:30335ms step_avg:42.43ms
step:716/1575 train_time:30393ms step_avg:42.45ms
step:717/1575 train_time:30456ms step_avg:42.48ms
step:718/1575 train_time:30515ms step_avg:42.50ms
step:719/1575 train_time:30578ms step_avg:42.53ms
step:720/1575 train_time:30638ms step_avg:42.55ms
step:721/1575 train_time:30701ms step_avg:42.58ms
step:722/1575 train_time:30761ms step_avg:42.60ms
step:723/1575 train_time:30823ms step_avg:42.63ms
step:724/1575 train_time:30883ms step_avg:42.66ms
step:725/1575 train_time:30947ms step_avg:42.69ms
step:726/1575 train_time:31006ms step_avg:42.71ms
step:727/1575 train_time:31069ms step_avg:42.74ms
step:728/1575 train_time:31128ms step_avg:42.76ms
step:729/1575 train_time:31192ms step_avg:42.79ms
step:730/1575 train_time:31251ms step_avg:42.81ms
step:731/1575 train_time:31315ms step_avg:42.84ms
step:732/1575 train_time:31375ms step_avg:42.86ms
step:733/1575 train_time:31437ms step_avg:42.89ms
step:734/1575 train_time:31497ms step_avg:42.91ms
step:735/1575 train_time:31560ms step_avg:42.94ms
step:736/1575 train_time:31619ms step_avg:42.96ms
step:737/1575 train_time:31682ms step_avg:42.99ms
step:738/1575 train_time:31742ms step_avg:43.01ms
step:739/1575 train_time:31804ms step_avg:43.04ms
step:740/1575 train_time:31864ms step_avg:43.06ms
step:741/1575 train_time:31928ms step_avg:43.09ms
step:742/1575 train_time:31987ms step_avg:43.11ms
step:743/1575 train_time:32051ms step_avg:43.14ms
step:744/1575 train_time:32110ms step_avg:43.16ms
step:745/1575 train_time:32174ms step_avg:43.19ms
step:746/1575 train_time:32233ms step_avg:43.21ms
step:747/1575 train_time:32297ms step_avg:43.24ms
step:748/1575 train_time:32356ms step_avg:43.26ms
step:749/1575 train_time:32419ms step_avg:43.28ms
step:750/1575 train_time:32478ms step_avg:43.30ms
step:750/1575 val_loss:3.8868 train_time:32526ms step_avg:43.37ms
step:751/1575 train_time:32546ms step_avg:43.34ms
step:752/1575 train_time:32606ms step_avg:43.36ms
step:753/1575 train_time:32671ms step_avg:43.39ms
step:754/1575 train_time:32732ms step_avg:43.41ms
step:755/1575 train_time:32796ms step_avg:43.44ms
step:756/1575 train_time:32855ms step_avg:43.46ms
step:757/1575 train_time:32918ms step_avg:43.48ms
step:758/1575 train_time:32977ms step_avg:43.51ms
step:759/1575 train_time:33039ms step_avg:43.53ms
step:760/1575 train_time:33098ms step_avg:43.55ms
step:761/1575 train_time:33161ms step_avg:43.58ms
step:762/1575 train_time:33220ms step_avg:43.60ms
step:763/1575 train_time:33282ms step_avg:43.62ms
step:764/1575 train_time:33341ms step_avg:43.64ms
step:765/1575 train_time:33404ms step_avg:43.67ms
step:766/1575 train_time:33465ms step_avg:43.69ms
step:767/1575 train_time:33529ms step_avg:43.72ms
step:768/1575 train_time:33589ms step_avg:43.74ms
step:769/1575 train_time:33653ms step_avg:43.76ms
step:770/1575 train_time:33713ms step_avg:43.78ms
step:771/1575 train_time:33777ms step_avg:43.81ms
step:772/1575 train_time:33837ms step_avg:43.83ms
step:773/1575 train_time:33900ms step_avg:43.86ms
step:774/1575 train_time:33960ms step_avg:43.88ms
step:775/1575 train_time:34022ms step_avg:43.90ms
step:776/1575 train_time:34081ms step_avg:43.92ms
step:777/1575 train_time:34144ms step_avg:43.94ms
step:778/1575 train_time:34203ms step_avg:43.96ms
step:779/1575 train_time:34266ms step_avg:43.99ms
step:780/1575 train_time:34324ms step_avg:44.01ms
step:781/1575 train_time:34387ms step_avg:44.03ms
step:782/1575 train_time:34446ms step_avg:44.05ms
step:783/1575 train_time:34510ms step_avg:44.07ms
step:784/1575 train_time:34569ms step_avg:44.09ms
step:785/1575 train_time:34635ms step_avg:44.12ms
step:786/1575 train_time:34697ms step_avg:44.14ms
step:787/1575 train_time:34759ms step_avg:44.17ms
step:788/1575 train_time:34819ms step_avg:44.19ms
step:789/1575 train_time:34882ms step_avg:44.21ms
step:790/1575 train_time:34941ms step_avg:44.23ms
step:791/1575 train_time:35004ms step_avg:44.25ms
step:792/1575 train_time:35063ms step_avg:44.27ms
step:793/1575 train_time:35126ms step_avg:44.30ms
step:794/1575 train_time:35185ms step_avg:44.31ms
step:795/1575 train_time:35248ms step_avg:44.34ms
step:796/1575 train_time:35307ms step_avg:44.36ms
step:797/1575 train_time:35370ms step_avg:44.38ms
step:798/1575 train_time:35430ms step_avg:44.40ms
step:799/1575 train_time:35493ms step_avg:44.42ms
step:800/1575 train_time:35552ms step_avg:44.44ms
step:801/1575 train_time:35616ms step_avg:44.46ms
step:802/1575 train_time:35675ms step_avg:44.48ms
step:803/1575 train_time:35739ms step_avg:44.51ms
step:804/1575 train_time:35798ms step_avg:44.53ms
step:805/1575 train_time:35862ms step_avg:44.55ms
step:806/1575 train_time:35922ms step_avg:44.57ms
step:807/1575 train_time:35986ms step_avg:44.59ms
step:808/1575 train_time:36044ms step_avg:44.61ms
step:809/1575 train_time:36107ms step_avg:44.63ms
step:810/1575 train_time:36167ms step_avg:44.65ms
step:811/1575 train_time:36230ms step_avg:44.67ms
step:812/1575 train_time:36289ms step_avg:44.69ms
step:813/1575 train_time:36351ms step_avg:44.71ms
step:814/1575 train_time:36411ms step_avg:44.73ms
step:815/1575 train_time:36475ms step_avg:44.75ms
step:816/1575 train_time:36534ms step_avg:44.77ms
step:817/1575 train_time:36598ms step_avg:44.80ms
step:818/1575 train_time:36656ms step_avg:44.81ms
step:819/1575 train_time:36721ms step_avg:44.84ms
step:820/1575 train_time:36780ms step_avg:44.85ms
step:821/1575 train_time:36843ms step_avg:44.88ms
step:822/1575 train_time:36902ms step_avg:44.89ms
step:823/1575 train_time:36966ms step_avg:44.92ms
step:824/1575 train_time:37024ms step_avg:44.93ms
step:825/1575 train_time:37088ms step_avg:44.95ms
step:826/1575 train_time:37147ms step_avg:44.97ms
step:827/1575 train_time:37210ms step_avg:44.99ms
step:828/1575 train_time:37269ms step_avg:45.01ms
step:829/1575 train_time:37332ms step_avg:45.03ms
step:830/1575 train_time:37392ms step_avg:45.05ms
step:831/1575 train_time:37455ms step_avg:45.07ms
step:832/1575 train_time:37515ms step_avg:45.09ms
step:833/1575 train_time:37578ms step_avg:45.11ms
step:834/1575 train_time:37638ms step_avg:45.13ms
step:835/1575 train_time:37701ms step_avg:45.15ms
step:836/1575 train_time:37760ms step_avg:45.17ms
step:837/1575 train_time:37824ms step_avg:45.19ms
step:838/1575 train_time:37883ms step_avg:45.21ms
step:839/1575 train_time:37947ms step_avg:45.23ms
step:840/1575 train_time:38006ms step_avg:45.24ms
step:841/1575 train_time:38068ms step_avg:45.27ms
step:842/1575 train_time:38128ms step_avg:45.28ms
step:843/1575 train_time:38191ms step_avg:45.30ms
step:844/1575 train_time:38250ms step_avg:45.32ms
step:845/1575 train_time:38313ms step_avg:45.34ms
step:846/1575 train_time:38372ms step_avg:45.36ms
step:847/1575 train_time:38436ms step_avg:45.38ms
step:848/1575 train_time:38495ms step_avg:45.40ms
step:849/1575 train_time:38558ms step_avg:45.42ms
step:850/1575 train_time:38618ms step_avg:45.43ms
step:851/1575 train_time:38682ms step_avg:45.45ms
step:852/1575 train_time:38742ms step_avg:45.47ms
step:853/1575 train_time:38805ms step_avg:45.49ms
step:854/1575 train_time:38865ms step_avg:45.51ms
step:855/1575 train_time:38928ms step_avg:45.53ms
step:856/1575 train_time:38987ms step_avg:45.55ms
step:857/1575 train_time:39050ms step_avg:45.57ms
step:858/1575 train_time:39110ms step_avg:45.58ms
step:859/1575 train_time:39172ms step_avg:45.60ms
step:860/1575 train_time:39232ms step_avg:45.62ms
step:861/1575 train_time:39295ms step_avg:45.64ms
step:862/1575 train_time:39354ms step_avg:45.65ms
step:863/1575 train_time:39417ms step_avg:45.67ms
step:864/1575 train_time:39477ms step_avg:45.69ms
step:865/1575 train_time:39543ms step_avg:45.71ms
step:866/1575 train_time:39600ms step_avg:45.73ms
step:867/1575 train_time:39663ms step_avg:45.75ms
step:868/1575 train_time:39723ms step_avg:45.76ms
step:869/1575 train_time:39786ms step_avg:45.78ms
step:870/1575 train_time:39845ms step_avg:45.80ms
step:871/1575 train_time:39908ms step_avg:45.82ms
step:872/1575 train_time:39967ms step_avg:45.83ms
step:873/1575 train_time:40030ms step_avg:45.85ms
step:874/1575 train_time:40090ms step_avg:45.87ms
step:875/1575 train_time:40154ms step_avg:45.89ms
step:876/1575 train_time:40212ms step_avg:45.90ms
step:877/1575 train_time:40276ms step_avg:45.92ms
step:878/1575 train_time:40335ms step_avg:45.94ms
step:879/1575 train_time:40399ms step_avg:45.96ms
step:880/1575 train_time:40458ms step_avg:45.98ms
step:881/1575 train_time:40521ms step_avg:45.99ms
step:882/1575 train_time:40581ms step_avg:46.01ms
step:883/1575 train_time:40644ms step_avg:46.03ms
step:884/1575 train_time:40704ms step_avg:46.04ms
step:885/1575 train_time:40767ms step_avg:46.06ms
step:886/1575 train_time:40832ms step_avg:46.09ms
step:887/1575 train_time:40893ms step_avg:46.10ms
step:888/1575 train_time:40952ms step_avg:46.12ms
step:889/1575 train_time:41015ms step_avg:46.14ms
step:890/1575 train_time:41075ms step_avg:46.15ms
step:891/1575 train_time:41137ms step_avg:46.17ms
step:892/1575 train_time:41196ms step_avg:46.18ms
step:893/1575 train_time:41260ms step_avg:46.20ms
step:894/1575 train_time:41320ms step_avg:46.22ms
step:895/1575 train_time:41382ms step_avg:46.24ms
step:896/1575 train_time:41442ms step_avg:46.25ms
step:897/1575 train_time:41506ms step_avg:46.27ms
step:898/1575 train_time:41564ms step_avg:46.28ms
step:899/1575 train_time:41627ms step_avg:46.30ms
step:900/1575 train_time:41686ms step_avg:46.32ms
step:901/1575 train_time:41750ms step_avg:46.34ms
step:902/1575 train_time:41809ms step_avg:46.35ms
step:903/1575 train_time:41873ms step_avg:46.37ms
step:904/1575 train_time:41932ms step_avg:46.38ms
step:905/1575 train_time:41995ms step_avg:46.40ms
step:906/1575 train_time:42055ms step_avg:46.42ms
step:907/1575 train_time:42119ms step_avg:46.44ms
step:908/1575 train_time:42178ms step_avg:46.45ms
step:909/1575 train_time:42242ms step_avg:46.47ms
step:910/1575 train_time:42302ms step_avg:46.49ms
step:911/1575 train_time:42365ms step_avg:46.50ms
step:912/1575 train_time:42424ms step_avg:46.52ms
step:913/1575 train_time:42487ms step_avg:46.54ms
step:914/1575 train_time:42546ms step_avg:46.55ms
step:915/1575 train_time:42610ms step_avg:46.57ms
step:916/1575 train_time:42670ms step_avg:46.58ms
step:917/1575 train_time:42733ms step_avg:46.60ms
step:918/1575 train_time:42793ms step_avg:46.62ms
step:919/1575 train_time:42856ms step_avg:46.63ms
step:920/1575 train_time:42915ms step_avg:46.65ms
step:921/1575 train_time:42978ms step_avg:46.66ms
step:922/1575 train_time:43037ms step_avg:46.68ms
step:923/1575 train_time:43100ms step_avg:46.70ms
step:924/1575 train_time:43161ms step_avg:46.71ms
step:925/1575 train_time:43223ms step_avg:46.73ms
step:926/1575 train_time:43282ms step_avg:46.74ms
step:927/1575 train_time:43345ms step_avg:46.76ms
step:928/1575 train_time:43404ms step_avg:46.77ms
step:929/1575 train_time:43468ms step_avg:46.79ms
step:930/1575 train_time:43527ms step_avg:46.80ms
step:931/1575 train_time:43591ms step_avg:46.82ms
step:932/1575 train_time:43650ms step_avg:46.83ms
step:933/1575 train_time:43714ms step_avg:46.85ms
step:934/1575 train_time:43774ms step_avg:46.87ms
step:935/1575 train_time:43838ms step_avg:46.89ms
step:936/1575 train_time:43897ms step_avg:46.90ms
step:937/1575 train_time:43960ms step_avg:46.92ms
step:938/1575 train_time:44019ms step_avg:46.93ms
step:939/1575 train_time:44082ms step_avg:46.95ms
step:940/1575 train_time:44141ms step_avg:46.96ms
step:941/1575 train_time:44204ms step_avg:46.98ms
step:942/1575 train_time:44263ms step_avg:46.99ms
step:943/1575 train_time:44326ms step_avg:47.01ms
step:944/1575 train_time:44385ms step_avg:47.02ms
step:945/1575 train_time:44449ms step_avg:47.04ms
step:946/1575 train_time:44508ms step_avg:47.05ms
step:947/1575 train_time:44571ms step_avg:47.07ms
step:948/1575 train_time:44631ms step_avg:47.08ms
step:949/1575 train_time:44694ms step_avg:47.10ms
step:950/1575 train_time:44753ms step_avg:47.11ms
step:951/1575 train_time:44816ms step_avg:47.13ms
step:952/1575 train_time:44875ms step_avg:47.14ms
step:953/1575 train_time:44939ms step_avg:47.16ms
step:954/1575 train_time:44998ms step_avg:47.17ms
step:955/1575 train_time:45061ms step_avg:47.18ms
step:956/1575 train_time:45121ms step_avg:47.20ms
step:957/1575 train_time:45184ms step_avg:47.21ms
step:958/1575 train_time:45243ms step_avg:47.23ms
step:959/1575 train_time:45306ms step_avg:47.24ms
step:960/1575 train_time:45365ms step_avg:47.26ms
step:961/1575 train_time:45429ms step_avg:47.27ms
step:962/1575 train_time:45488ms step_avg:47.28ms
step:963/1575 train_time:45551ms step_avg:47.30ms
step:964/1575 train_time:45610ms step_avg:47.31ms
step:965/1575 train_time:45673ms step_avg:47.33ms
step:966/1575 train_time:45733ms step_avg:47.34ms
step:967/1575 train_time:45796ms step_avg:47.36ms
step:968/1575 train_time:45856ms step_avg:47.37ms
step:969/1575 train_time:45920ms step_avg:47.39ms
step:970/1575 train_time:45980ms step_avg:47.40ms
step:971/1575 train_time:46043ms step_avg:47.42ms
step:972/1575 train_time:46102ms step_avg:47.43ms
step:973/1575 train_time:46166ms step_avg:47.45ms
step:974/1575 train_time:46225ms step_avg:47.46ms
step:975/1575 train_time:46288ms step_avg:47.47ms
step:976/1575 train_time:46347ms step_avg:47.49ms
step:977/1575 train_time:46410ms step_avg:47.50ms
step:978/1575 train_time:46470ms step_avg:47.51ms
step:979/1575 train_time:46532ms step_avg:47.53ms
step:980/1575 train_time:46592ms step_avg:47.54ms
step:981/1575 train_time:46655ms step_avg:47.56ms
step:982/1575 train_time:46715ms step_avg:47.57ms
step:983/1575 train_time:46778ms step_avg:47.59ms
step:984/1575 train_time:46838ms step_avg:47.60ms
step:985/1575 train_time:46901ms step_avg:47.62ms
step:986/1575 train_time:46961ms step_avg:47.63ms
step:987/1575 train_time:47024ms step_avg:47.64ms
step:988/1575 train_time:47083ms step_avg:47.65ms
step:989/1575 train_time:47147ms step_avg:47.67ms
step:990/1575 train_time:47206ms step_avg:47.68ms
step:991/1575 train_time:47268ms step_avg:47.70ms
step:992/1575 train_time:47328ms step_avg:47.71ms
step:993/1575 train_time:47390ms step_avg:47.72ms
step:994/1575 train_time:47449ms step_avg:47.74ms
step:995/1575 train_time:47512ms step_avg:47.75ms
step:996/1575 train_time:47572ms step_avg:47.76ms
step:997/1575 train_time:47636ms step_avg:47.78ms
step:998/1575 train_time:47696ms step_avg:47.79ms
step:999/1575 train_time:47759ms step_avg:47.81ms
step:1000/1575 train_time:47818ms step_avg:47.82ms
step:1000/1575 val_loss:3.5833 train_time:47866ms step_avg:47.87ms
step:1001/1575 train_time:47886ms step_avg:47.84ms
step:1002/1575 train_time:47944ms step_avg:47.85ms
step:1003/1575 train_time:48012ms step_avg:47.87ms
step:1004/1575 train_time:48074ms step_avg:47.88ms
step:1005/1575 train_time:48137ms step_avg:47.90ms
step:1006/1575 train_time:48197ms step_avg:47.91ms
step:1007/1575 train_time:48260ms step_avg:47.92ms
step:1008/1575 train_time:48319ms step_avg:47.94ms
step:1009/1575 train_time:48382ms step_avg:47.95ms
step:1010/1575 train_time:48441ms step_avg:47.96ms
step:1011/1575 train_time:48504ms step_avg:47.98ms
step:1012/1575 train_time:48563ms step_avg:47.99ms
step:1013/1575 train_time:48625ms step_avg:48.00ms
step:1014/1575 train_time:48684ms step_avg:48.01ms
step:1015/1575 train_time:48747ms step_avg:48.03ms
step:1016/1575 train_time:48809ms step_avg:48.04ms
step:1017/1575 train_time:48871ms step_avg:48.05ms
step:1018/1575 train_time:48931ms step_avg:48.07ms
step:1019/1575 train_time:48996ms step_avg:48.08ms
step:1020/1575 train_time:49055ms step_avg:48.09ms
step:1021/1575 train_time:49119ms step_avg:48.11ms
step:1022/1575 train_time:49178ms step_avg:48.12ms
step:1023/1575 train_time:49242ms step_avg:48.14ms
step:1024/1575 train_time:49303ms step_avg:48.15ms
step:1025/1575 train_time:49375ms step_avg:48.17ms
step:1026/1575 train_time:49458ms step_avg:48.20ms
step:1027/1575 train_time:49549ms step_avg:48.25ms
step:1028/1575 train_time:49633ms step_avg:48.28ms
step:1029/1575 train_time:49721ms step_avg:48.32ms
step:1030/1575 train_time:49806ms step_avg:48.36ms
step:1031/1575 train_time:49897ms step_avg:48.40ms
step:1032/1575 train_time:49983ms step_avg:48.43ms
step:1033/1575 train_time:50074ms step_avg:48.47ms
step:1034/1575 train_time:50161ms step_avg:48.51ms
step:1035/1575 train_time:50251ms step_avg:48.55ms
step:1036/1575 train_time:50336ms step_avg:48.59ms
step:1037/1575 train_time:50425ms step_avg:48.63ms
step:1038/1575 train_time:50511ms step_avg:48.66ms
step:1039/1575 train_time:50600ms step_avg:48.70ms
step:1040/1575 train_time:50685ms step_avg:48.74ms
step:1041/1575 train_time:50775ms step_avg:48.77ms
step:1042/1575 train_time:50860ms step_avg:48.81ms
step:1043/1575 train_time:50949ms step_avg:48.85ms
step:1044/1575 train_time:51036ms step_avg:48.88ms
step:1045/1575 train_time:51127ms step_avg:48.93ms
step:1046/1575 train_time:51213ms step_avg:48.96ms
step:1047/1575 train_time:51303ms step_avg:49.00ms
step:1048/1575 train_time:51388ms step_avg:49.03ms
step:1049/1575 train_time:51477ms step_avg:49.07ms
step:1050/1575 train_time:51562ms step_avg:49.11ms
step:1051/1575 train_time:51652ms step_avg:49.15ms
step:1052/1575 train_time:51738ms step_avg:49.18ms
step:1053/1575 train_time:51827ms step_avg:49.22ms
step:1054/1575 train_time:51913ms step_avg:49.25ms
step:1055/1575 train_time:52003ms step_avg:49.29ms
step:1056/1575 train_time:52089ms step_avg:49.33ms
step:1057/1575 train_time:52179ms step_avg:49.37ms
step:1058/1575 train_time:52265ms step_avg:49.40ms
step:1059/1575 train_time:52354ms step_avg:49.44ms
step:1060/1575 train_time:52440ms step_avg:49.47ms
step:1061/1575 train_time:52528ms step_avg:49.51ms
step:1062/1575 train_time:52615ms step_avg:49.54ms
step:1063/1575 train_time:52704ms step_avg:49.58ms
step:1064/1575 train_time:52790ms step_avg:49.61ms
step:1065/1575 train_time:52879ms step_avg:49.65ms
step:1066/1575 train_time:52965ms step_avg:49.69ms
step:1067/1575 train_time:53056ms step_avg:49.72ms
step:1068/1575 train_time:53141ms step_avg:49.76ms
step:1069/1575 train_time:53230ms step_avg:49.79ms
step:1070/1575 train_time:53316ms step_avg:49.83ms
step:1071/1575 train_time:53406ms step_avg:49.87ms
step:1072/1575 train_time:53492ms step_avg:49.90ms
step:1073/1575 train_time:53581ms step_avg:49.94ms
step:1074/1575 train_time:53670ms step_avg:49.97ms
step:1075/1575 train_time:53758ms step_avg:50.01ms
step:1076/1575 train_time:53842ms step_avg:50.04ms
step:1077/1575 train_time:53932ms step_avg:50.08ms
step:1078/1575 train_time:54017ms step_avg:50.11ms
step:1079/1575 train_time:54106ms step_avg:50.14ms
step:1080/1575 train_time:54191ms step_avg:50.18ms
step:1081/1575 train_time:54282ms step_avg:50.21ms
step:1082/1575 train_time:54367ms step_avg:50.25ms
step:1083/1575 train_time:54457ms step_avg:50.28ms
step:1084/1575 train_time:54543ms step_avg:50.32ms
step:1085/1575 train_time:54632ms step_avg:50.35ms
step:1086/1575 train_time:54717ms step_avg:50.38ms
step:1087/1575 train_time:54806ms step_avg:50.42ms
step:1088/1575 train_time:54893ms step_avg:50.45ms
step:1089/1575 train_time:54983ms step_avg:50.49ms
step:1090/1575 train_time:55068ms step_avg:50.52ms
step:1091/1575 train_time:55159ms step_avg:50.56ms
step:1092/1575 train_time:55245ms step_avg:50.59ms
step:1093/1575 train_time:55333ms step_avg:50.63ms
step:1094/1575 train_time:55419ms step_avg:50.66ms
step:1095/1575 train_time:55509ms step_avg:50.69ms
step:1096/1575 train_time:55595ms step_avg:50.72ms
step:1097/1575 train_time:55684ms step_avg:50.76ms
step:1098/1575 train_time:55770ms step_avg:50.79ms
step:1099/1575 train_time:55859ms step_avg:50.83ms
step:1100/1575 train_time:55945ms step_avg:50.86ms
step:1101/1575 train_time:56035ms step_avg:50.90ms
step:1102/1575 train_time:56121ms step_avg:50.93ms
step:1103/1575 train_time:56209ms step_avg:50.96ms
step:1104/1575 train_time:56296ms step_avg:50.99ms
step:1105/1575 train_time:56386ms step_avg:51.03ms
step:1106/1575 train_time:56471ms step_avg:51.06ms
step:1107/1575 train_time:56561ms step_avg:51.09ms
step:1108/1575 train_time:56647ms step_avg:51.13ms
step:1109/1575 train_time:56735ms step_avg:51.16ms
step:1110/1575 train_time:56822ms step_avg:51.19ms
step:1111/1575 train_time:56911ms step_avg:51.22ms
step:1112/1575 train_time:56997ms step_avg:51.26ms
step:1113/1575 train_time:57088ms step_avg:51.29ms
step:1114/1575 train_time:57171ms step_avg:51.32ms
step:1115/1575 train_time:57262ms step_avg:51.36ms
step:1116/1575 train_time:57357ms step_avg:51.40ms
step:1117/1575 train_time:57442ms step_avg:51.42ms
step:1118/1575 train_time:57527ms step_avg:51.46ms
step:1119/1575 train_time:57616ms step_avg:51.49ms
step:1120/1575 train_time:57703ms step_avg:51.52ms
step:1121/1575 train_time:57792ms step_avg:51.55ms
step:1122/1575 train_time:57877ms step_avg:51.58ms
step:1123/1575 train_time:57967ms step_avg:51.62ms
step:1124/1575 train_time:58053ms step_avg:51.65ms
step:1125/1575 train_time:58142ms step_avg:51.68ms
step:1126/1575 train_time:58228ms step_avg:51.71ms
step:1127/1575 train_time:58316ms step_avg:51.74ms
step:1128/1575 train_time:58407ms step_avg:51.78ms
step:1129/1575 train_time:58491ms step_avg:51.81ms
step:1130/1575 train_time:58577ms step_avg:51.84ms
step:1131/1575 train_time:58668ms step_avg:51.87ms
step:1132/1575 train_time:58754ms step_avg:51.90ms
step:1133/1575 train_time:58843ms step_avg:51.94ms
step:1134/1575 train_time:58928ms step_avg:51.96ms
step:1135/1575 train_time:59017ms step_avg:52.00ms
step:1136/1575 train_time:59104ms step_avg:52.03ms
step:1137/1575 train_time:59193ms step_avg:52.06ms
step:1138/1575 train_time:59278ms step_avg:52.09ms
step:1139/1575 train_time:59368ms step_avg:52.12ms
step:1140/1575 train_time:59453ms step_avg:52.15ms
step:1141/1575 train_time:59543ms step_avg:52.18ms
step:1142/1575 train_time:59629ms step_avg:52.21ms
step:1143/1575 train_time:59718ms step_avg:52.25ms
step:1144/1575 train_time:59804ms step_avg:52.28ms
step:1145/1575 train_time:59893ms step_avg:52.31ms
step:1146/1575 train_time:59984ms step_avg:52.34ms
step:1147/1575 train_time:60066ms step_avg:52.37ms
step:1148/1575 train_time:60156ms step_avg:52.40ms
step:1149/1575 train_time:60245ms step_avg:52.43ms
step:1150/1575 train_time:60331ms step_avg:52.46ms
step:1151/1575 train_time:60420ms step_avg:52.49ms
step:1152/1575 train_time:60505ms step_avg:52.52ms
step:1153/1575 train_time:60595ms step_avg:52.55ms
step:1154/1575 train_time:60680ms step_avg:52.58ms
step:1155/1575 train_time:60765ms step_avg:52.61ms
step:1156/1575 train_time:60856ms step_avg:52.64ms
step:1157/1575 train_time:60946ms step_avg:52.68ms
step:1158/1575 train_time:61032ms step_avg:52.70ms
step:1159/1575 train_time:61122ms step_avg:52.74ms
step:1160/1575 train_time:61207ms step_avg:52.76ms
step:1161/1575 train_time:61296ms step_avg:52.80ms
step:1162/1575 train_time:61380ms step_avg:52.82ms
step:1163/1575 train_time:61471ms step_avg:52.86ms
step:1164/1575 train_time:61556ms step_avg:52.88ms
step:1165/1575 train_time:61646ms step_avg:52.91ms
step:1166/1575 train_time:61727ms step_avg:52.94ms
step:1167/1575 train_time:61820ms step_avg:52.97ms
step:1168/1575 train_time:61907ms step_avg:53.00ms
step:1169/1575 train_time:61996ms step_avg:53.03ms
step:1170/1575 train_time:62081ms step_avg:53.06ms
step:1171/1575 train_time:62172ms step_avg:53.09ms
step:1172/1575 train_time:62259ms step_avg:53.12ms
step:1173/1575 train_time:62349ms step_avg:53.15ms
step:1174/1575 train_time:62433ms step_avg:53.18ms
step:1175/1575 train_time:62524ms step_avg:53.21ms
step:1176/1575 train_time:62609ms step_avg:53.24ms
step:1177/1575 train_time:62694ms step_avg:53.27ms
step:1178/1575 train_time:62785ms step_avg:53.30ms
step:1179/1575 train_time:62874ms step_avg:53.33ms
step:1180/1575 train_time:62959ms step_avg:53.36ms
step:1181/1575 train_time:63049ms step_avg:53.39ms
step:1182/1575 train_time:63134ms step_avg:53.41ms
step:1183/1575 train_time:63224ms step_avg:53.44ms
step:1184/1575 train_time:63309ms step_avg:53.47ms
step:1185/1575 train_time:63398ms step_avg:53.50ms
step:1186/1575 train_time:63484ms step_avg:53.53ms
step:1187/1575 train_time:63574ms step_avg:53.56ms
step:1188/1575 train_time:63660ms step_avg:53.59ms
step:1189/1575 train_time:63753ms step_avg:53.62ms
step:1190/1575 train_time:63840ms step_avg:53.65ms
step:1191/1575 train_time:63925ms step_avg:53.67ms
step:1192/1575 train_time:64010ms step_avg:53.70ms
step:1193/1575 train_time:64095ms step_avg:53.73ms
step:1194/1575 train_time:64186ms step_avg:53.76ms
step:1195/1575 train_time:64275ms step_avg:53.79ms
step:1196/1575 train_time:64361ms step_avg:53.81ms
step:1197/1575 train_time:64450ms step_avg:53.84ms
step:1198/1575 train_time:64536ms step_avg:53.87ms
step:1199/1575 train_time:64625ms step_avg:53.90ms
step:1200/1575 train_time:64711ms step_avg:53.93ms
step:1201/1575 train_time:64800ms step_avg:53.96ms
step:1202/1575 train_time:64886ms step_avg:53.98ms
step:1203/1575 train_time:64976ms step_avg:54.01ms
step:1204/1575 train_time:65062ms step_avg:54.04ms
step:1205/1575 train_time:65151ms step_avg:54.07ms
step:1206/1575 train_time:65237ms step_avg:54.09ms
step:1207/1575 train_time:65327ms step_avg:54.12ms
step:1208/1575 train_time:65411ms step_avg:54.15ms
step:1209/1575 train_time:65501ms step_avg:54.18ms
step:1210/1575 train_time:65587ms step_avg:54.20ms
step:1211/1575 train_time:65676ms step_avg:54.23ms
step:1212/1575 train_time:65762ms step_avg:54.26ms
step:1213/1575 train_time:65852ms step_avg:54.29ms
step:1214/1575 train_time:65937ms step_avg:54.31ms
step:1215/1575 train_time:66022ms step_avg:54.34ms
step:1216/1575 train_time:66112ms step_avg:54.37ms
step:1217/1575 train_time:66204ms step_avg:54.40ms
step:1218/1575 train_time:66288ms step_avg:54.42ms
step:1219/1575 train_time:66377ms step_avg:54.45ms
step:1220/1575 train_time:66463ms step_avg:54.48ms
step:1221/1575 train_time:66552ms step_avg:54.51ms
step:1222/1575 train_time:66638ms step_avg:54.53ms
step:1223/1575 train_time:66729ms step_avg:54.56ms
step:1224/1575 train_time:66814ms step_avg:54.59ms
step:1225/1575 train_time:66903ms step_avg:54.61ms
step:1226/1575 train_time:66984ms step_avg:54.64ms
step:1227/1575 train_time:67078ms step_avg:54.67ms
step:1228/1575 train_time:67163ms step_avg:54.69ms
step:1229/1575 train_time:67252ms step_avg:54.72ms
step:1230/1575 train_time:67334ms step_avg:54.74ms
step:1231/1575 train_time:67422ms step_avg:54.77ms
step:1232/1575 train_time:67508ms step_avg:54.80ms
step:1233/1575 train_time:67598ms step_avg:54.82ms
step:1234/1575 train_time:67684ms step_avg:54.85ms
step:1235/1575 train_time:67774ms step_avg:54.88ms
step:1236/1575 train_time:67859ms step_avg:54.90ms
step:1237/1575 train_time:67948ms step_avg:54.93ms
step:1238/1575 train_time:68034ms step_avg:54.95ms
step:1239/1575 train_time:68123ms step_avg:54.98ms
step:1240/1575 train_time:68209ms step_avg:55.01ms
step:1241/1575 train_time:68301ms step_avg:55.04ms
step:1242/1575 train_time:68387ms step_avg:55.06ms
step:1243/1575 train_time:68476ms step_avg:55.09ms
step:1244/1575 train_time:68562ms step_avg:55.11ms
step:1245/1575 train_time:68654ms step_avg:55.14ms
step:1246/1575 train_time:68739ms step_avg:55.17ms
step:1247/1575 train_time:68827ms step_avg:55.19ms
step:1248/1575 train_time:68913ms step_avg:55.22ms
step:1249/1575 train_time:69002ms step_avg:55.25ms
step:1250/1575 train_time:69088ms step_avg:55.27ms
step:1250/1575 val_loss:3.4071 train_time:69163ms step_avg:55.33ms
step:1251/1575 train_time:69182ms step_avg:55.30ms
step:1252/1575 train_time:69269ms step_avg:55.33ms
step:1253/1575 train_time:69364ms step_avg:55.36ms
step:1254/1575 train_time:69452ms step_avg:55.38ms
step:1255/1575 train_time:69541ms step_avg:55.41ms
step:1256/1575 train_time:69625ms step_avg:55.43ms
step:1257/1575 train_time:69713ms step_avg:55.46ms
step:1258/1575 train_time:69798ms step_avg:55.48ms
step:1259/1575 train_time:69886ms step_avg:55.51ms
step:1260/1575 train_time:69971ms step_avg:55.53ms
step:1261/1575 train_time:70059ms step_avg:55.56ms
step:1262/1575 train_time:70145ms step_avg:55.58ms
step:1263/1575 train_time:70237ms step_avg:55.61ms
step:1264/1575 train_time:70325ms step_avg:55.64ms
step:1265/1575 train_time:70415ms step_avg:55.66ms
step:1266/1575 train_time:70501ms step_avg:55.69ms
step:1267/1575 train_time:70590ms step_avg:55.71ms
step:1268/1575 train_time:70675ms step_avg:55.74ms
step:1269/1575 train_time:70764ms step_avg:55.76ms
step:1270/1575 train_time:70848ms step_avg:55.79ms
step:1271/1575 train_time:70936ms step_avg:55.81ms
step:1272/1575 train_time:71021ms step_avg:55.83ms
step:1273/1575 train_time:71109ms step_avg:55.86ms
step:1274/1575 train_time:71196ms step_avg:55.88ms
step:1275/1575 train_time:71287ms step_avg:55.91ms
step:1276/1575 train_time:71374ms step_avg:55.94ms
step:1277/1575 train_time:71465ms step_avg:55.96ms
step:1278/1575 train_time:71552ms step_avg:55.99ms
step:1279/1575 train_time:71641ms step_avg:56.01ms
step:1280/1575 train_time:71725ms step_avg:56.04ms
step:1281/1575 train_time:71814ms step_avg:56.06ms
step:1282/1575 train_time:71899ms step_avg:56.08ms
step:1283/1575 train_time:71988ms step_avg:56.11ms
step:1284/1575 train_time:72073ms step_avg:56.13ms
step:1285/1575 train_time:72163ms step_avg:56.16ms
step:1286/1575 train_time:72251ms step_avg:56.18ms
step:1287/1575 train_time:72344ms step_avg:56.21ms
step:1288/1575 train_time:72429ms step_avg:56.23ms
step:1289/1575 train_time:72521ms step_avg:56.26ms
step:1290/1575 train_time:72606ms step_avg:56.28ms
step:1291/1575 train_time:72695ms step_avg:56.31ms
step:1292/1575 train_time:72780ms step_avg:56.33ms
step:1293/1575 train_time:72868ms step_avg:56.36ms
step:1294/1575 train_time:72953ms step_avg:56.38ms
step:1295/1575 train_time:73042ms step_avg:56.40ms
step:1296/1575 train_time:73127ms step_avg:56.43ms
step:1297/1575 train_time:73217ms step_avg:56.45ms
step:1298/1575 train_time:73304ms step_avg:56.47ms
step:1299/1575 train_time:73395ms step_avg:56.50ms
step:1300/1575 train_time:73481ms step_avg:56.52ms
step:1301/1575 train_time:73571ms step_avg:56.55ms
step:1302/1575 train_time:73656ms step_avg:56.57ms
step:1303/1575 train_time:73745ms step_avg:56.60ms
step:1304/1575 train_time:73830ms step_avg:56.62ms
step:1305/1575 train_time:73921ms step_avg:56.64ms
step:1306/1575 train_time:74006ms step_avg:56.67ms
step:1307/1575 train_time:74095ms step_avg:56.69ms
step:1308/1575 train_time:74180ms step_avg:56.71ms
step:1309/1575 train_time:74270ms step_avg:56.74ms
step:1310/1575 train_time:74356ms step_avg:56.76ms
step:1311/1575 train_time:74446ms step_avg:56.79ms
step:1312/1575 train_time:74532ms step_avg:56.81ms
step:1313/1575 train_time:74623ms step_avg:56.83ms
step:1314/1575 train_time:74707ms step_avg:56.85ms
step:1315/1575 train_time:74796ms step_avg:56.88ms
step:1316/1575 train_time:74882ms step_avg:56.90ms
step:1317/1575 train_time:74970ms step_avg:56.92ms
step:1318/1575 train_time:75056ms step_avg:56.95ms
step:1319/1575 train_time:75146ms step_avg:56.97ms
step:1320/1575 train_time:75231ms step_avg:56.99ms
step:1321/1575 train_time:75322ms step_avg:57.02ms
step:1322/1575 train_time:75407ms step_avg:57.04ms
step:1323/1575 train_time:75497ms step_avg:57.07ms
step:1324/1575 train_time:75583ms step_avg:57.09ms
step:1325/1575 train_time:75672ms step_avg:57.11ms
step:1326/1575 train_time:75759ms step_avg:57.13ms
step:1327/1575 train_time:75848ms step_avg:57.16ms
step:1328/1575 train_time:75933ms step_avg:57.18ms
step:1329/1575 train_time:76022ms step_avg:57.20ms
step:1330/1575 train_time:76107ms step_avg:57.22ms
step:1331/1575 train_time:76199ms step_avg:57.25ms
step:1332/1575 train_time:76284ms step_avg:57.27ms
step:1333/1575 train_time:76373ms step_avg:57.29ms
step:1334/1575 train_time:76459ms step_avg:57.32ms
step:1335/1575 train_time:76549ms step_avg:57.34ms
step:1336/1575 train_time:76635ms step_avg:57.36ms
step:1337/1575 train_time:76724ms step_avg:57.38ms
step:1338/1575 train_time:76809ms step_avg:57.41ms
step:1339/1575 train_time:76899ms step_avg:57.43ms
step:1340/1575 train_time:76983ms step_avg:57.45ms
step:1341/1575 train_time:77073ms step_avg:57.47ms
step:1342/1575 train_time:77159ms step_avg:57.50ms
step:1343/1575 train_time:77250ms step_avg:57.52ms
step:1344/1575 train_time:77336ms step_avg:57.54ms
step:1345/1575 train_time:77425ms step_avg:57.57ms
step:1346/1575 train_time:77511ms step_avg:57.59ms
step:1347/1575 train_time:77601ms step_avg:57.61ms
step:1348/1575 train_time:77687ms step_avg:57.63ms
step:1349/1575 train_time:77776ms step_avg:57.65ms
step:1350/1575 train_time:77862ms step_avg:57.68ms
step:1351/1575 train_time:77954ms step_avg:57.70ms
step:1352/1575 train_time:78038ms step_avg:57.72ms
step:1353/1575 train_time:78126ms step_avg:57.74ms
step:1354/1575 train_time:78212ms step_avg:57.76ms
step:1355/1575 train_time:78301ms step_avg:57.79ms
step:1356/1575 train_time:78386ms step_avg:57.81ms
step:1357/1575 train_time:78476ms step_avg:57.83ms
step:1358/1575 train_time:78566ms step_avg:57.85ms
step:1359/1575 train_time:78655ms step_avg:57.88ms
step:1360/1575 train_time:78740ms step_avg:57.90ms
step:1361/1575 train_time:78829ms step_avg:57.92ms
step:1362/1575 train_time:78915ms step_avg:57.94ms
step:1363/1575 train_time:79005ms step_avg:57.96ms
step:1364/1575 train_time:79090ms step_avg:57.98ms
step:1365/1575 train_time:79180ms step_avg:58.01ms
step:1366/1575 train_time:79266ms step_avg:58.03ms
step:1367/1575 train_time:79355ms step_avg:58.05ms
step:1368/1575 train_time:79442ms step_avg:58.07ms
step:1369/1575 train_time:79533ms step_avg:58.10ms
step:1370/1575 train_time:79618ms step_avg:58.12ms
step:1371/1575 train_time:79708ms step_avg:58.14ms
step:1372/1575 train_time:79793ms step_avg:58.16ms
step:1373/1575 train_time:79883ms step_avg:58.18ms
step:1374/1575 train_time:79969ms step_avg:58.20ms
step:1375/1575 train_time:80057ms step_avg:58.22ms
step:1376/1575 train_time:80144ms step_avg:58.24ms
step:1377/1575 train_time:80233ms step_avg:58.27ms
step:1378/1575 train_time:80321ms step_avg:58.29ms
step:1379/1575 train_time:80408ms step_avg:58.31ms
step:1380/1575 train_time:80495ms step_avg:58.33ms
step:1381/1575 train_time:80585ms step_avg:58.35ms
step:1382/1575 train_time:80671ms step_avg:58.37ms
step:1383/1575 train_time:80761ms step_avg:58.40ms
step:1384/1575 train_time:80846ms step_avg:58.41ms
step:1385/1575 train_time:80936ms step_avg:58.44ms
step:1386/1575 train_time:81020ms step_avg:58.46ms
step:1387/1575 train_time:81110ms step_avg:58.48ms
step:1388/1575 train_time:81196ms step_avg:58.50ms
step:1389/1575 train_time:81286ms step_avg:58.52ms
step:1390/1575 train_time:81371ms step_avg:58.54ms
step:1391/1575 train_time:81460ms step_avg:58.56ms
step:1392/1575 train_time:81546ms step_avg:58.58ms
step:1393/1575 train_time:81636ms step_avg:58.60ms
step:1394/1575 train_time:81721ms step_avg:58.62ms
step:1395/1575 train_time:81810ms step_avg:58.65ms
step:1396/1575 train_time:81896ms step_avg:58.66ms
step:1397/1575 train_time:81985ms step_avg:58.69ms
step:1398/1575 train_time:82072ms step_avg:58.71ms
step:1399/1575 train_time:82161ms step_avg:58.73ms
step:1400/1575 train_time:82247ms step_avg:58.75ms
step:1401/1575 train_time:82337ms step_avg:58.77ms
step:1402/1575 train_time:82422ms step_avg:58.79ms
step:1403/1575 train_time:82513ms step_avg:58.81ms
step:1404/1575 train_time:82598ms step_avg:58.83ms
step:1405/1575 train_time:82688ms step_avg:58.85ms
step:1406/1575 train_time:82777ms step_avg:58.87ms
step:1407/1575 train_time:82865ms step_avg:58.89ms
step:1408/1575 train_time:82949ms step_avg:58.91ms
step:1409/1575 train_time:83038ms step_avg:58.93ms
step:1410/1575 train_time:83123ms step_avg:58.95ms
step:1411/1575 train_time:83213ms step_avg:58.97ms
step:1412/1575 train_time:83299ms step_avg:58.99ms
step:1413/1575 train_time:83388ms step_avg:59.02ms
step:1414/1575 train_time:83475ms step_avg:59.03ms
step:1415/1575 train_time:83564ms step_avg:59.06ms
step:1416/1575 train_time:83650ms step_avg:59.07ms
step:1417/1575 train_time:83740ms step_avg:59.10ms
step:1418/1575 train_time:83825ms step_avg:59.11ms
step:1419/1575 train_time:83915ms step_avg:59.14ms
step:1420/1575 train_time:84000ms step_avg:59.16ms
step:1421/1575 train_time:84089ms step_avg:59.18ms
step:1422/1575 train_time:84176ms step_avg:59.20ms
step:1423/1575 train_time:84265ms step_avg:59.22ms
step:1424/1575 train_time:84353ms step_avg:59.24ms
step:1425/1575 train_time:84442ms step_avg:59.26ms
step:1426/1575 train_time:84527ms step_avg:59.28ms
step:1427/1575 train_time:84618ms step_avg:59.30ms
step:1428/1575 train_time:84702ms step_avg:59.32ms
step:1429/1575 train_time:84791ms step_avg:59.34ms
step:1430/1575 train_time:84877ms step_avg:59.35ms
step:1431/1575 train_time:84967ms step_avg:59.38ms
step:1432/1575 train_time:85052ms step_avg:59.39ms
step:1433/1575 train_time:85142ms step_avg:59.41ms
step:1434/1575 train_time:85227ms step_avg:59.43ms
step:1435/1575 train_time:85317ms step_avg:59.45ms
step:1436/1575 train_time:85402ms step_avg:59.47ms
step:1437/1575 train_time:85492ms step_avg:59.49ms
step:1438/1575 train_time:85579ms step_avg:59.51ms
step:1439/1575 train_time:85668ms step_avg:59.53ms
step:1440/1575 train_time:85754ms step_avg:59.55ms
step:1441/1575 train_time:85843ms step_avg:59.57ms
step:1442/1575 train_time:85929ms step_avg:59.59ms
step:1443/1575 train_time:86018ms step_avg:59.61ms
step:1444/1575 train_time:86103ms step_avg:59.63ms
step:1445/1575 train_time:86193ms step_avg:59.65ms
step:1446/1575 train_time:86279ms step_avg:59.67ms
step:1447/1575 train_time:86368ms step_avg:59.69ms
step:1448/1575 train_time:86454ms step_avg:59.71ms
step:1449/1575 train_time:86544ms step_avg:59.73ms
step:1450/1575 train_time:86630ms step_avg:59.74ms
step:1451/1575 train_time:86720ms step_avg:59.77ms
step:1452/1575 train_time:86805ms step_avg:59.78ms
step:1453/1575 train_time:86894ms step_avg:59.80ms
step:1454/1575 train_time:86980ms step_avg:59.82ms
step:1455/1575 train_time:87070ms step_avg:59.84ms
step:1456/1575 train_time:87156ms step_avg:59.86ms
step:1457/1575 train_time:87245ms step_avg:59.88ms
step:1458/1575 train_time:87330ms step_avg:59.90ms
step:1459/1575 train_time:87421ms step_avg:59.92ms
step:1460/1575 train_time:87506ms step_avg:59.94ms
step:1461/1575 train_time:87595ms step_avg:59.96ms
step:1462/1575 train_time:87681ms step_avg:59.97ms
step:1463/1575 train_time:87770ms step_avg:59.99ms
step:1464/1575 train_time:87856ms step_avg:60.01ms
step:1465/1575 train_time:87945ms step_avg:60.03ms
step:1466/1575 train_time:88030ms step_avg:60.05ms
step:1467/1575 train_time:88121ms step_avg:60.07ms
step:1468/1575 train_time:88206ms step_avg:60.09ms
step:1469/1575 train_time:88296ms step_avg:60.11ms
step:1470/1575 train_time:88383ms step_avg:60.12ms
step:1471/1575 train_time:88472ms step_avg:60.14ms
step:1472/1575 train_time:88558ms step_avg:60.16ms
step:1473/1575 train_time:88648ms step_avg:60.18ms
step:1474/1575 train_time:88736ms step_avg:60.20ms
step:1475/1575 train_time:88825ms step_avg:60.22ms
step:1476/1575 train_time:88912ms step_avg:60.24ms
step:1477/1575 train_time:89001ms step_avg:60.26ms
step:1478/1575 train_time:89087ms step_avg:60.28ms
step:1479/1575 train_time:89176ms step_avg:60.30ms
step:1480/1575 train_time:89261ms step_avg:60.31ms
step:1481/1575 train_time:89351ms step_avg:60.33ms
step:1482/1575 train_time:89437ms step_avg:60.35ms
step:1483/1575 train_time:89527ms step_avg:60.37ms
step:1484/1575 train_time:89613ms step_avg:60.39ms
step:1485/1575 train_time:89703ms step_avg:60.41ms
step:1486/1575 train_time:89788ms step_avg:60.42ms
step:1487/1575 train_time:89878ms step_avg:60.44ms
step:1488/1575 train_time:89964ms step_avg:60.46ms
step:1489/1575 train_time:90053ms step_avg:60.48ms
step:1490/1575 train_time:90139ms step_avg:60.50ms
step:1491/1575 train_time:90227ms step_avg:60.51ms
step:1492/1575 train_time:90313ms step_avg:60.53ms
step:1493/1575 train_time:90402ms step_avg:60.55ms
step:1494/1575 train_time:90488ms step_avg:60.57ms
step:1495/1575 train_time:90577ms step_avg:60.59ms
step:1496/1575 train_time:90663ms step_avg:60.60ms
step:1497/1575 train_time:90753ms step_avg:60.62ms
step:1498/1575 train_time:90838ms step_avg:60.64ms
step:1499/1575 train_time:90928ms step_avg:60.66ms
step:1500/1575 train_time:91013ms step_avg:60.68ms
step:1500/1575 val_loss:3.3010 train_time:91087ms step_avg:60.72ms
step:1501/1575 train_time:91106ms step_avg:60.70ms
step:1502/1575 train_time:91192ms step_avg:60.71ms
step:1503/1575 train_time:91284ms step_avg:60.73ms
step:1504/1575 train_time:91369ms step_avg:60.75ms
step:1505/1575 train_time:91459ms step_avg:60.77ms
step:1506/1575 train_time:91544ms step_avg:60.79ms
step:1507/1575 train_time:91632ms step_avg:60.80ms
step:1508/1575 train_time:91717ms step_avg:60.82ms
step:1509/1575 train_time:91805ms step_avg:60.84ms
step:1510/1575 train_time:91891ms step_avg:60.86ms
step:1511/1575 train_time:91980ms step_avg:60.87ms
step:1512/1575 train_time:92066ms step_avg:60.89ms
step:1513/1575 train_time:92162ms step_avg:60.91ms
step:1514/1575 train_time:92250ms step_avg:60.93ms
step:1515/1575 train_time:92341ms step_avg:60.95ms
step:1516/1575 train_time:92430ms step_avg:60.97ms
step:1517/1575 train_time:92517ms step_avg:60.99ms
step:1518/1575 train_time:92601ms step_avg:61.00ms
step:1519/1575 train_time:92689ms step_avg:61.02ms
step:1520/1575 train_time:92774ms step_avg:61.04ms
step:1521/1575 train_time:92863ms step_avg:61.05ms
step:1522/1575 train_time:92948ms step_avg:61.07ms
step:1523/1575 train_time:93038ms step_avg:61.09ms
step:1524/1575 train_time:93125ms step_avg:61.11ms
step:1525/1575 train_time:93217ms step_avg:61.13ms
step:1526/1575 train_time:93304ms step_avg:61.14ms
step:1527/1575 train_time:93393ms step_avg:61.16ms
step:1528/1575 train_time:93479ms step_avg:61.18ms
step:1529/1575 train_time:93567ms step_avg:61.20ms
step:1530/1575 train_time:93653ms step_avg:61.21ms
step:1531/1575 train_time:93742ms step_avg:61.23ms
step:1532/1575 train_time:93827ms step_avg:61.24ms
step:1533/1575 train_time:93916ms step_avg:61.26ms
step:1534/1575 train_time:94002ms step_avg:61.28ms
step:1535/1575 train_time:94093ms step_avg:61.30ms
step:1536/1575 train_time:94186ms step_avg:61.32ms
step:1537/1575 train_time:94275ms step_avg:61.34ms
step:1538/1575 train_time:94362ms step_avg:61.35ms
step:1539/1575 train_time:94452ms step_avg:61.37ms
step:1540/1575 train_time:94538ms step_avg:61.39ms
step:1541/1575 train_time:94627ms step_avg:61.41ms
step:1542/1575 train_time:94713ms step_avg:61.42ms
step:1543/1575 train_time:94803ms step_avg:61.44ms
step:1544/1575 train_time:94888ms step_avg:61.46ms
step:1545/1575 train_time:94979ms step_avg:61.47ms
step:1546/1575 train_time:95066ms step_avg:61.49ms
step:1547/1575 train_time:95158ms step_avg:61.51ms
step:1548/1575 train_time:95244ms step_avg:61.53ms
step:1549/1575 train_time:95333ms step_avg:61.55ms
step:1550/1575 train_time:95420ms step_avg:61.56ms
step:1551/1575 train_time:95510ms step_avg:61.58ms
step:1552/1575 train_time:95596ms step_avg:61.60ms
step:1553/1575 train_time:95685ms step_avg:61.61ms
step:1554/1575 train_time:95771ms step_avg:61.63ms
step:1555/1575 train_time:95860ms step_avg:61.65ms
step:1556/1575 train_time:95946ms step_avg:61.66ms
step:1557/1575 train_time:96035ms step_avg:61.68ms
step:1558/1575 train_time:96122ms step_avg:61.70ms
step:1559/1575 train_time:96213ms step_avg:61.71ms
step:1560/1575 train_time:96300ms step_avg:61.73ms
step:1561/1575 train_time:96391ms step_avg:61.75ms
step:1562/1575 train_time:96477ms step_avg:61.76ms
step:1563/1575 train_time:96569ms step_avg:61.78ms
step:1564/1575 train_time:96654ms step_avg:61.80ms
step:1565/1575 train_time:96743ms step_avg:61.82ms
step:1566/1575 train_time:96830ms step_avg:61.83ms
step:1567/1575 train_time:96920ms step_avg:61.85ms
step:1568/1575 train_time:97005ms step_avg:61.87ms
step:1569/1575 train_time:97098ms step_avg:61.89ms
step:1570/1575 train_time:97182ms step_avg:61.90ms
step:1571/1575 train_time:97271ms step_avg:61.92ms
step:1572/1575 train_time:97357ms step_avg:61.93ms
step:1573/1575 train_time:97446ms step_avg:61.95ms
step:1574/1575 train_time:97533ms step_avg:61.96ms
step:1575/1575 train_time:97622ms step_avg:61.98ms
step:1575/1575 val_loss:3.2794 train_time:97691ms step_avg:62.03ms
peak memory allocated: 30933 MiB reserved: 46998 MiB