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:16:48 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   34C    P0            120W /  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   40C    P0            122W /  700W |    1519MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   3  NVIDIA H100 80GB HBM3          On  |   00000000:64:00.0 Off |                    0 |
| N/A   34C    P0            120W /  700W |    1519MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   4  NVIDIA H100 80GB HBM3          On  |   00000000:6A:00.0 Off |                    0 |
| N/A   34C    P0            120W /  700W |    1519MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   5  NVIDIA H100 80GB HBM3          On  |   00000000:6B:00.0 Off |                    0 |
| N/A   41C    P0            129W /  700W |    1519MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   6  NVIDIA H100 80GB HBM3          On  |   00000000:6C:00.0 Off |                    0 |
| N/A   39C    P0            124W /  700W |    1519MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   7  NVIDIA H100 80GB HBM3          On  |   00000000:6D:00.0 Off |                    0 |
| N/A   35C    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          237197      C   /usr/bin/python3                       1510MiB |
|    1   N/A  N/A          237198      C   /usr/bin/python3                       1510MiB |
|    2   N/A  N/A          237199      C   /usr/bin/python3                       1510MiB |
|    3   N/A  N/A          237200      C   /usr/bin/python3                       1510MiB |
|    4   N/A  N/A          237201      C   /usr/bin/python3                       1510MiB |
|    5   N/A  N/A          237202      C   /usr/bin/python3                       1510MiB |
|    6   N/A  N/A          237203      C   /usr/bin/python3                       1510MiB |
|    7   N/A  N/A          237204      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.8310 train_time:0ms step_avg:0.03ms
step:1/1575 train_time:83ms step_avg:83.13ms
step:2/1575 train_time:105ms step_avg:52.28ms
step:3/1575 train_time:123ms step_avg:40.87ms
step:4/1575 train_time:155ms step_avg:38.66ms
step:5/1575 train_time:185ms step_avg:37.07ms
step:6/1575 train_time:257ms step_avg:42.88ms
step:7/1575 train_time:274ms step_avg:39.20ms
step:8/1575 train_time:330ms step_avg:41.30ms
step:9/1575 train_time:361ms step_avg:40.15ms
step:10/1575 train_time:400ms step_avg:39.98ms
step:11/1575 train_time:431ms step_avg:39.17ms
step:12/1575 train_time:470ms step_avg:39.13ms
step:13/1575 train_time:501ms step_avg:38.51ms
step:14/1575 train_time:539ms step_avg:38.50ms
step:15/1575 train_time:571ms step_avg:38.06ms
step:16/1575 train_time:610ms step_avg:38.11ms
step:17/1575 train_time:641ms step_avg:37.69ms
step:18/1575 train_time:680ms step_avg:37.77ms
step:19/1575 train_time:711ms step_avg:37.41ms
step:20/1575 train_time:749ms step_avg:37.47ms
step:21/1575 train_time:781ms step_avg:37.19ms
step:22/1575 train_time:820ms step_avg:37.28ms
step:23/1575 train_time:851ms step_avg:37.02ms
step:24/1575 train_time:890ms step_avg:37.08ms
step:25/1575 train_time:921ms step_avg:36.83ms
step:26/1575 train_time:959ms step_avg:36.90ms
step:27/1575 train_time:991ms step_avg:36.71ms
step:28/1575 train_time:1030ms step_avg:36.79ms
step:29/1575 train_time:1061ms step_avg:36.58ms
step:30/1575 train_time:1099ms step_avg:36.65ms
step:31/1575 train_time:1131ms step_avg:36.47ms
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:1239ms step_avg:36.44ms
step:35/1575 train_time:1271ms step_avg:36.30ms
step:36/1575 train_time:1310ms step_avg:36.39ms
step:37/1575 train_time:1340ms step_avg:36.23ms
step:38/1575 train_time:1380ms step_avg:36.32ms
step:39/1575 train_time:1412ms step_avg:36.20ms
step:40/1575 train_time:1450ms step_avg:36.26ms
step:41/1575 train_time:1482ms step_avg:36.14ms
step:42/1575 train_time:1520ms step_avg:36.19ms
step:43/1575 train_time:1551ms step_avg:36.08ms
step:44/1575 train_time:1590ms step_avg:36.13ms
step:45/1575 train_time:1621ms step_avg:36.02ms
step:46/1575 train_time:1659ms step_avg:36.07ms
step:47/1575 train_time:1690ms step_avg:35.97ms
step:48/1575 train_time:1729ms step_avg:36.02ms
step:49/1575 train_time:1760ms step_avg:35.92ms
step:50/1575 train_time:1799ms step_avg:35.97ms
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.88ms
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.76ms
step:58/1575 train_time:2077ms step_avg:35.80ms
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.74ms
step:63/1575 train_time:2247ms 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:2387ms step_avg:35.63ms
step:68/1575 train_time:2426ms step_avg:35.67ms
step:69/1575 train_time:2457ms step_avg:35.61ms
step:70/1575 train_time:2495ms step_avg:35.64ms
step:71/1575 train_time:2527ms step_avg:35.59ms
step:72/1575 train_time:2566ms step_avg:35.64ms
step:73/1575 train_time:2597ms step_avg:35.57ms
step:74/1575 train_time:2635ms step_avg:35.61ms
step:75/1575 train_time:2666ms step_avg:35.55ms
step:76/1575 train_time:2705ms step_avg:35.60ms
step:77/1575 train_time:2737ms step_avg:35.54ms
step:78/1575 train_time:2775ms step_avg:35.58ms
step:79/1575 train_time:2806ms step_avg:35.52ms
step:80/1575 train_time:2844ms step_avg:35.55ms
step:81/1575 train_time:2875ms step_avg:35.50ms
step:82/1575 train_time:2914ms step_avg:35.54ms
step:83/1575 train_time:2945ms step_avg:35.48ms
step:84/1575 train_time:2984ms step_avg:35.52ms
step:85/1575 train_time:3014ms step_avg:35.46ms
step:86/1575 train_time:3053ms step_avg:35.50ms
step:87/1575 train_time:3084ms step_avg:35.45ms
step:88/1575 train_time:3123ms step_avg:35.49ms
step:89/1575 train_time:3154ms step_avg:35.44ms
step:90/1575 train_time:3192ms step_avg:35.47ms
step:91/1575 train_time:3223ms step_avg:35.42ms
step:92/1575 train_time:3262ms step_avg:35.46ms
step:93/1575 train_time:3293ms step_avg:35.41ms
step:94/1575 train_time:3332ms step_avg:35.45ms
step:95/1575 train_time:3363ms step_avg:35.40ms
step:96/1575 train_time:3402ms step_avg:35.43ms
step:97/1575 train_time:3433ms step_avg:35.39ms
step:98/1575 train_time:3471ms step_avg:35.42ms
step:99/1575 train_time:3503ms step_avg:35.38ms
step:100/1575 train_time:3541ms step_avg:35.41ms
step:101/1575 train_time:3572ms step_avg:35.37ms
step:102/1575 train_time:3611ms step_avg:35.40ms
step:103/1575 train_time:3642ms step_avg:35.36ms
step:104/1575 train_time:3680ms step_avg:35.39ms
step:105/1575 train_time:3712ms step_avg:35.35ms
step:106/1575 train_time:3750ms step_avg:35.37ms
step:107/1575 train_time:3781ms step_avg:35.33ms
step:108/1575 train_time:3820ms step_avg:35.37ms
step:109/1575 train_time:3851ms step_avg:35.33ms
step:110/1575 train_time:3890ms step_avg:35.36ms
step:111/1575 train_time:3921ms step_avg:35.32ms
step:112/1575 train_time:3959ms step_avg:35.35ms
step:113/1575 train_time:3990ms step_avg:35.31ms
step:114/1575 train_time:4030ms step_avg:35.35ms
step:115/1575 train_time:4061ms step_avg:35.31ms
step:116/1575 train_time:4100ms step_avg:35.34ms
step:117/1575 train_time:4131ms step_avg:35.31ms
step:118/1575 train_time:4169ms step_avg:35.33ms
step:119/1575 train_time:4200ms step_avg:35.30ms
step:120/1575 train_time:4240ms step_avg:35.33ms
step:121/1575 train_time:4271ms step_avg:35.30ms
step:122/1575 train_time:4310ms step_avg:35.33ms
step:123/1575 train_time:4341ms step_avg:35.29ms
step:124/1575 train_time:4379ms step_avg:35.32ms
step:125/1575 train_time:4410ms step_avg:35.28ms
step:126/1575 train_time:4449ms step_avg:35.31ms
step:127/1575 train_time:4480ms step_avg:35.28ms
step:128/1575 train_time:4519ms step_avg:35.31ms
step:129/1575 train_time:4550ms step_avg:35.27ms
step:130/1575 train_time:4589ms step_avg:35.30ms
step:131/1575 train_time:4620ms step_avg:35.27ms
step:132/1575 train_time:4658ms step_avg:35.29ms
step:133/1575 train_time:4689ms step_avg:35.26ms
step:134/1575 train_time:4728ms step_avg:35.28ms
step:135/1575 train_time:4759ms step_avg:35.25ms
step:136/1575 train_time:4798ms step_avg:35.28ms
step:137/1575 train_time:4828ms step_avg:35.24ms
step:138/1575 train_time:4867ms step_avg:35.27ms
step:139/1575 train_time:4898ms step_avg:35.24ms
step:140/1575 train_time:4936ms step_avg:35.26ms
step:141/1575 train_time:4967ms step_avg:35.23ms
step:142/1575 train_time:5006ms step_avg:35.26ms
step:143/1575 train_time:5037ms step_avg:35.23ms
step:144/1575 train_time:5076ms step_avg:35.25ms
step:145/1575 train_time:5106ms step_avg:35.22ms
step:146/1575 train_time:5145ms step_avg:35.24ms
step:147/1575 train_time:5176ms step_avg:35.21ms
step:148/1575 train_time:5214ms step_avg:35.23ms
step:149/1575 train_time:5246ms step_avg:35.21ms
step:150/1575 train_time:5285ms step_avg:35.23ms
step:151/1575 train_time:5316ms step_avg:35.21ms
step:152/1575 train_time:5354ms step_avg:35.22ms
step:153/1575 train_time:5385ms step_avg:35.20ms
step:154/1575 train_time:5423ms step_avg:35.22ms
step:155/1575 train_time:5454ms step_avg:35.19ms
step:156/1575 train_time:5493ms step_avg:35.21ms
step:157/1575 train_time:5524ms step_avg:35.18ms
step:158/1575 train_time:5563ms step_avg:35.21ms
step:159/1575 train_time:5593ms step_avg:35.18ms
step:160/1575 train_time:5632ms step_avg:35.20ms
step:161/1575 train_time:5663ms step_avg:35.17ms
step:162/1575 train_time:5701ms step_avg:35.19ms
step:163/1575 train_time:5732ms step_avg:35.17ms
step:164/1575 train_time:5770ms step_avg:35.18ms
step:165/1575 train_time:5801ms step_avg:35.16ms
step:166/1575 train_time:5840ms step_avg:35.18ms
step:167/1575 train_time:5871ms step_avg:35.16ms
step:168/1575 train_time:5910ms step_avg:35.18ms
step:169/1575 train_time:5941ms step_avg:35.15ms
step:170/1575 train_time:5979ms step_avg:35.17ms
step:171/1575 train_time:6011ms step_avg:35.15ms
step:172/1575 train_time:6049ms step_avg:35.17ms
step:173/1575 train_time:6080ms step_avg:35.15ms
step:174/1575 train_time:6119ms step_avg:35.17ms
step:175/1575 train_time:6150ms step_avg:35.14ms
step:176/1575 train_time:6188ms step_avg:35.16ms
step:177/1575 train_time:6219ms step_avg:35.14ms
step:178/1575 train_time:6258ms step_avg:35.16ms
step:179/1575 train_time:6289ms step_avg:35.13ms
step:180/1575 train_time:6328ms step_avg:35.15ms
step:181/1575 train_time:6359ms step_avg:35.13ms
step:182/1575 train_time:6398ms step_avg:35.15ms
step:183/1575 train_time:6429ms step_avg:35.13ms
step:184/1575 train_time:6467ms step_avg:35.15ms
step:185/1575 train_time:6498ms step_avg:35.13ms
step:186/1575 train_time:6536ms step_avg:35.14ms
step:187/1575 train_time:6567ms step_avg:35.12ms
step:188/1575 train_time:6606ms step_avg:35.14ms
step:189/1575 train_time:6637ms step_avg:35.12ms
step:190/1575 train_time:6675ms step_avg:35.13ms
step:191/1575 train_time:6707ms step_avg:35.11ms
step:192/1575 train_time:6746ms step_avg:35.13ms
step:193/1575 train_time:6777ms step_avg:35.11ms
step:194/1575 train_time:6817ms step_avg:35.14ms
step:195/1575 train_time:6846ms step_avg:35.11ms
step:196/1575 train_time:6885ms step_avg:35.13ms
step:197/1575 train_time:6915ms step_avg:35.10ms
step:198/1575 train_time:6954ms step_avg:35.12ms
step:199/1575 train_time:6985ms step_avg:35.10ms
step:200/1575 train_time:7024ms step_avg:35.12ms
step:201/1575 train_time:7055ms step_avg:35.10ms
step:202/1575 train_time:7093ms step_avg:35.11ms
step:203/1575 train_time:7124ms step_avg:35.09ms
step:204/1575 train_time:7163ms step_avg:35.11ms
step:205/1575 train_time:7194ms step_avg:35.09ms
step:206/1575 train_time:7232ms step_avg:35.11ms
step:207/1575 train_time:7263ms step_avg:35.09ms
step:208/1575 train_time:7302ms step_avg:35.11ms
step:209/1575 train_time:7333ms step_avg:35.09ms
step:210/1575 train_time:7372ms step_avg:35.10ms
step:211/1575 train_time:7403ms step_avg:35.08ms
step:212/1575 train_time:7441ms step_avg:35.10ms
step:213/1575 train_time:7472ms step_avg:35.08ms
step:214/1575 train_time:7511ms step_avg:35.10ms
step:215/1575 train_time:7542ms step_avg:35.08ms
step:216/1575 train_time:7580ms step_avg:35.09ms
step:217/1575 train_time:7611ms step_avg:35.07ms
step:218/1575 train_time:7649ms step_avg:35.09ms
step:219/1575 train_time:7680ms step_avg:35.07ms
step:220/1575 train_time:7719ms step_avg:35.08ms
step:221/1575 train_time:7750ms step_avg:35.07ms
step:222/1575 train_time:7788ms step_avg:35.08ms
step:223/1575 train_time:7819ms step_avg:35.06ms
step:224/1575 train_time:7857ms step_avg:35.08ms
step:225/1575 train_time:7889ms step_avg:35.06ms
step:226/1575 train_time:7928ms step_avg:35.08ms
step:227/1575 train_time:7959ms step_avg:35.06ms
step:228/1575 train_time:7998ms step_avg:35.08ms
step:229/1575 train_time:8028ms step_avg:35.06ms
step:230/1575 train_time:8067ms step_avg:35.07ms
step:231/1575 train_time:8098ms step_avg:35.06ms
step:232/1575 train_time:8136ms step_avg:35.07ms
step:233/1575 train_time:8168ms step_avg:35.06ms
step:234/1575 train_time:8206ms step_avg:35.07ms
step:235/1575 train_time:8237ms step_avg:35.05ms
step:236/1575 train_time:8276ms step_avg:35.07ms
step:237/1575 train_time:8307ms step_avg:35.05ms
step:238/1575 train_time:8346ms step_avg:35.07ms
step:239/1575 train_time:8377ms step_avg:35.05ms
step:240/1575 train_time:8415ms step_avg:35.06ms
step:241/1575 train_time:8446ms step_avg:35.05ms
step:242/1575 train_time:8485ms step_avg:35.06ms
step:243/1575 train_time:8516ms step_avg:35.04ms
step:244/1575 train_time:8554ms step_avg:35.06ms
step:245/1575 train_time:8585ms step_avg:35.04ms
step:246/1575 train_time:8624ms step_avg:35.06ms
step:247/1575 train_time:8655ms step_avg:35.04ms
step:248/1575 train_time:8693ms step_avg:35.05ms
step:249/1575 train_time:8724ms step_avg:35.04ms
step:250/1575 train_time:8763ms step_avg:35.05ms
step:250/1575 val_loss:4.5744 train_time:8812ms step_avg:35.25ms
step:251/1575 train_time:8831ms step_avg:35.18ms
step:252/1575 train_time:8850ms step_avg:35.12ms
step:253/1575 train_time:8867ms step_avg:35.05ms
step:254/1575 train_time:8906ms step_avg:35.06ms
step:255/1575 train_time:8938ms step_avg:35.05ms
step:256/1575 train_time:8977ms step_avg:35.07ms
step:257/1575 train_time:9009ms step_avg:35.05ms
step:258/1575 train_time:9048ms step_avg:35.07ms
step:259/1575 train_time:9079ms step_avg:35.05ms
step:260/1575 train_time:9117ms step_avg:35.07ms
step:261/1575 train_time:9149ms step_avg:35.05ms
step:262/1575 train_time:9187ms step_avg:35.07ms
step:263/1575 train_time:9218ms step_avg:35.05ms
step:264/1575 train_time:9257ms step_avg:35.06ms
step:265/1575 train_time:9288ms step_avg:35.05ms
step:266/1575 train_time:9326ms step_avg:35.06ms
step:267/1575 train_time:9357ms step_avg:35.04ms
step:268/1575 train_time:9395ms step_avg:35.06ms
step:269/1575 train_time:9427ms step_avg:35.04ms
step:270/1575 train_time:9466ms step_avg:35.06ms
step:271/1575 train_time:9496ms step_avg:35.04ms
step:272/1575 train_time:9534ms step_avg:35.05ms
step:273/1575 train_time:9565ms step_avg:35.04ms
step:274/1575 train_time:9604ms step_avg:35.05ms
step:275/1575 train_time:9635ms step_avg:35.04ms
step:276/1575 train_time:9673ms step_avg:35.05ms
step:277/1575 train_time:9704ms step_avg:35.03ms
step:278/1575 train_time:9743ms step_avg:35.05ms
step:279/1575 train_time:9774ms step_avg:35.03ms
step:280/1575 train_time:9812ms step_avg:35.04ms
step:281/1575 train_time:9843ms step_avg:35.03ms
step:282/1575 train_time:9881ms step_avg:35.04ms
step:283/1575 train_time:9912ms step_avg:35.02ms
step:284/1575 train_time:9951ms step_avg:35.04ms
step:285/1575 train_time:9982ms step_avg:35.02ms
step:286/1575 train_time:10020ms step_avg:35.03ms
step:287/1575 train_time:10051ms step_avg:35.02ms
step:288/1575 train_time:10089ms step_avg:35.03ms
step:289/1575 train_time:10120ms step_avg:35.02ms
step:290/1575 train_time:10158ms step_avg:35.03ms
step:291/1575 train_time:10190ms step_avg:35.02ms
step:292/1575 train_time:10228ms step_avg:35.03ms
step:293/1575 train_time:10259ms step_avg:35.01ms
step:294/1575 train_time:10297ms step_avg:35.02ms
step:295/1575 train_time:10328ms step_avg:35.01ms
step:296/1575 train_time:10367ms step_avg:35.02ms
step:297/1575 train_time:10397ms step_avg:35.01ms
step:298/1575 train_time:10436ms step_avg:35.02ms
step:299/1575 train_time:10467ms step_avg:35.01ms
step:300/1575 train_time:10505ms step_avg:35.02ms
step:301/1575 train_time:10536ms step_avg:35.00ms
step:302/1575 train_time:10575ms step_avg:35.02ms
step:303/1575 train_time:10606ms step_avg:35.00ms
step:304/1575 train_time:10644ms step_avg:35.01ms
step:305/1575 train_time:10676ms step_avg:35.00ms
step:306/1575 train_time:10714ms step_avg:35.01ms
step:307/1575 train_time:10745ms step_avg:35.00ms
step:308/1575 train_time:10783ms step_avg:35.01ms
step:309/1575 train_time:10814ms step_avg:35.00ms
step:310/1575 train_time:10853ms step_avg:35.01ms
step:311/1575 train_time:10884ms step_avg:35.00ms
step:312/1575 train_time:10922ms step_avg:35.01ms
step:313/1575 train_time:10954ms step_avg:35.00ms
step:314/1575 train_time:10992ms step_avg:35.01ms
step:315/1575 train_time:11023ms step_avg:34.99ms
step:316/1575 train_time:11062ms step_avg:35.01ms
step:317/1575 train_time:11093ms step_avg:34.99ms
step:318/1575 train_time:11131ms step_avg:35.00ms
step:319/1575 train_time:11162ms step_avg:34.99ms
step:320/1575 train_time:11200ms step_avg:35.00ms
step:321/1575 train_time:11231ms step_avg:34.99ms
step:322/1575 train_time:11270ms step_avg:35.00ms
step:323/1575 train_time:11301ms step_avg:34.99ms
step:324/1575 train_time:11339ms step_avg:35.00ms
step:325/1575 train_time:11370ms step_avg:34.98ms
step:326/1575 train_time:11408ms step_avg:34.99ms
step:327/1575 train_time:11439ms step_avg:34.98ms
step:328/1575 train_time:11477ms step_avg:34.99ms
step:329/1575 train_time:11508ms step_avg:34.98ms
step:330/1575 train_time:11547ms step_avg:34.99ms
step:331/1575 train_time:11578ms step_avg:34.98ms
step:332/1575 train_time:11616ms step_avg:34.99ms
step:333/1575 train_time:11647ms step_avg:34.98ms
step:334/1575 train_time:11686ms step_avg:34.99ms
step:335/1575 train_time:11717ms step_avg:34.98ms
step:336/1575 train_time:11755ms step_avg:34.99ms
step:337/1575 train_time:11786ms step_avg:34.97ms
step:338/1575 train_time:11824ms step_avg:34.98ms
step:339/1575 train_time:11856ms step_avg:34.97ms
step:340/1575 train_time:11894ms step_avg:34.98ms
step:341/1575 train_time:11925ms step_avg:34.97ms
step:342/1575 train_time:11965ms step_avg:34.98ms
step:343/1575 train_time:11995ms step_avg:34.97ms
step:344/1575 train_time:12033ms step_avg:34.98ms
step:345/1575 train_time:12064ms step_avg:34.97ms
step:346/1575 train_time:12103ms step_avg:34.98ms
step:347/1575 train_time:12134ms step_avg:34.97ms
step:348/1575 train_time:12172ms step_avg:34.98ms
step:349/1575 train_time:12203ms step_avg:34.97ms
step:350/1575 train_time:12241ms step_avg:34.97ms
step:351/1575 train_time:12272ms step_avg:34.96ms
step:352/1575 train_time:12310ms step_avg:34.97ms
step:353/1575 train_time:12341ms step_avg:34.96ms
step:354/1575 train_time:12379ms step_avg:34.97ms
step:355/1575 train_time:12410ms step_avg:34.96ms
step:356/1575 train_time:12449ms step_avg:34.97ms
step:357/1575 train_time:12480ms step_avg:34.96ms
step:358/1575 train_time:12518ms step_avg:34.97ms
step:359/1575 train_time:12549ms step_avg:34.95ms
step:360/1575 train_time:12587ms step_avg:34.96ms
step:361/1575 train_time:12618ms step_avg:34.95ms
step:362/1575 train_time:12656ms step_avg:34.96ms
step:363/1575 train_time:12687ms step_avg:34.95ms
step:364/1575 train_time:12726ms step_avg:34.96ms
step:365/1575 train_time:12756ms step_avg:34.95ms
step:366/1575 train_time:12795ms step_avg:34.96ms
step:367/1575 train_time:12826ms step_avg:34.95ms
step:368/1575 train_time:12865ms step_avg:34.96ms
step:369/1575 train_time:12896ms step_avg:34.95ms
step:370/1575 train_time:12934ms step_avg:34.96ms
step:371/1575 train_time:12965ms step_avg:34.95ms
step:372/1575 train_time:13003ms step_avg:34.96ms
step:373/1575 train_time:13034ms step_avg:34.94ms
step:374/1575 train_time:13073ms step_avg:34.96ms
step:375/1575 train_time:13104ms step_avg:34.94ms
step:376/1575 train_time:13143ms step_avg:34.95ms
step:377/1575 train_time:13174ms step_avg:34.94ms
step:378/1575 train_time:13212ms step_avg:34.95ms
step:379/1575 train_time:13243ms step_avg:34.94ms
step:380/1575 train_time:13281ms step_avg:34.95ms
step:381/1575 train_time:13312ms step_avg:34.94ms
step:382/1575 train_time:13350ms step_avg:34.95ms
step:383/1575 train_time:13381ms step_avg:34.94ms
step:384/1575 train_time:13420ms step_avg:34.95ms
step:385/1575 train_time:13451ms step_avg:34.94ms
step:386/1575 train_time:13489ms step_avg:34.95ms
step:387/1575 train_time:13520ms step_avg:34.93ms
step:388/1575 train_time:13558ms step_avg:34.94ms
step:389/1575 train_time:13588ms step_avg:34.93ms
step:390/1575 train_time:13627ms 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:13765ms step_avg:34.94ms
step:395/1575 train_time:13796ms step_avg:34.93ms
step:396/1575 train_time:13834ms step_avg:34.94ms
step:397/1575 train_time:13865ms step_avg:34.93ms
step:398/1575 train_time:13904ms step_avg:34.93ms
step:399/1575 train_time:13935ms step_avg:34.92ms
step:400/1575 train_time:13974ms step_avg:34.93ms
step:401/1575 train_time:14004ms step_avg:34.92ms
step:402/1575 train_time:14042ms step_avg:34.93ms
step:403/1575 train_time:14073ms step_avg:34.92ms
step:404/1575 train_time:14112ms step_avg:34.93ms
step:405/1575 train_time:14143ms step_avg:34.92ms
step:406/1575 train_time:14181ms 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:14320ms step_avg:34.93ms
step:411/1575 train_time:14351ms 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:14459ms step_avg:34.92ms
step:415/1575 train_time:14490ms step_avg:34.92ms
step:416/1575 train_time:14528ms step_avg:34.92ms
step:417/1575 train_time:14559ms step_avg:34.91ms
step:418/1575 train_time:14597ms step_avg:34.92ms
step:419/1575 train_time:14629ms step_avg:34.91ms
step:420/1575 train_time:14667ms step_avg:34.92ms
step:421/1575 train_time:14698ms step_avg:34.91ms
step:422/1575 train_time:14736ms step_avg:34.92ms
step:423/1575 train_time:14767ms step_avg:34.91ms
step:424/1575 train_time:14806ms step_avg:34.92ms
step:425/1575 train_time:14837ms step_avg:34.91ms
step:426/1575 train_time:14875ms step_avg:34.92ms
step:427/1575 train_time:14906ms step_avg:34.91ms
step:428/1575 train_time:14946ms step_avg:34.92ms
step:429/1575 train_time:14975ms step_avg:34.91ms
step:430/1575 train_time:15013ms step_avg:34.91ms
step:431/1575 train_time:15044ms step_avg:34.91ms
step:432/1575 train_time:15082ms step_avg:34.91ms
step:433/1575 train_time:15113ms step_avg:34.90ms
step:434/1575 train_time:15152ms step_avg:34.91ms
step:435/1575 train_time:15183ms step_avg:34.90ms
step:436/1575 train_time:15221ms step_avg:34.91ms
step:437/1575 train_time:15252ms step_avg:34.90ms
step:438/1575 train_time:15291ms step_avg:34.91ms
step:439/1575 train_time:15322ms step_avg:34.90ms
step:440/1575 train_time:15361ms step_avg:34.91ms
step:441/1575 train_time:15392ms step_avg:34.90ms
step:442/1575 train_time:15430ms step_avg:34.91ms
step:443/1575 train_time:15461ms step_avg:34.90ms
step:444/1575 train_time:15500ms step_avg:34.91ms
step:445/1575 train_time:15531ms step_avg:34.90ms
step:446/1575 train_time:15569ms step_avg:34.91ms
step:447/1575 train_time:15600ms step_avg:34.90ms
step:448/1575 train_time:15638ms step_avg:34.91ms
step:449/1575 train_time:15670ms step_avg:34.90ms
step:450/1575 train_time:15708ms step_avg:34.91ms
step:451/1575 train_time:15739ms step_avg:34.90ms
step:452/1575 train_time:15777ms step_avg:34.91ms
step:453/1575 train_time:15808ms step_avg:34.90ms
step:454/1575 train_time:15847ms step_avg:34.90ms
step:455/1575 train_time:15878ms step_avg:34.90ms
step:456/1575 train_time:15916ms step_avg:34.90ms
step:457/1575 train_time:15947ms step_avg:34.90ms
step:458/1575 train_time:15986ms step_avg:34.90ms
step:459/1575 train_time:16017ms step_avg:34.90ms
step:460/1575 train_time:16055ms step_avg:34.90ms
step:461/1575 train_time:16086ms step_avg:34.89ms
step:462/1575 train_time:16124ms step_avg:34.90ms
step:463/1575 train_time:16155ms step_avg:34.89ms
step:464/1575 train_time:16194ms step_avg:34.90ms
step:465/1575 train_time:16225ms step_avg:34.89ms
step:466/1575 train_time:16264ms step_avg:34.90ms
step:467/1575 train_time:16294ms step_avg:34.89ms
step:468/1575 train_time:16332ms step_avg:34.90ms
step:469/1575 train_time:16364ms step_avg:34.89ms
step:470/1575 train_time:16402ms step_avg:34.90ms
step:471/1575 train_time:16433ms step_avg:34.89ms
step:472/1575 train_time:16473ms step_avg:34.90ms
step:473/1575 train_time:16503ms step_avg:34.89ms
step:474/1575 train_time:16541ms step_avg:34.90ms
step:475/1575 train_time:16573ms step_avg:34.89ms
step:476/1575 train_time:16611ms step_avg:34.90ms
step:477/1575 train_time:16642ms step_avg:34.89ms
step:478/1575 train_time:16680ms step_avg:34.90ms
step:479/1575 train_time:16711ms step_avg:34.89ms
step:480/1575 train_time:16750ms step_avg:34.90ms
step:481/1575 train_time:16781ms step_avg:34.89ms
step:482/1575 train_time:16819ms step_avg:34.89ms
step:483/1575 train_time:16850ms step_avg:34.89ms
step:484/1575 train_time:16889ms step_avg:34.90ms
step:485/1575 train_time:16920ms step_avg:34.89ms
step:486/1575 train_time:16958ms step_avg:34.89ms
step:487/1575 train_time:16989ms step_avg:34.89ms
step:488/1575 train_time:17027ms step_avg:34.89ms
step:489/1575 train_time:17058ms step_avg:34.88ms
step:490/1575 train_time:17096ms step_avg:34.89ms
step:491/1575 train_time:17128ms step_avg:34.88ms
step:492/1575 train_time:17167ms step_avg:34.89ms
step:493/1575 train_time:17198ms step_avg:34.88ms
step:494/1575 train_time:17236ms step_avg:34.89ms
step:495/1575 train_time:17267ms step_avg:34.88ms
step:496/1575 train_time:17305ms step_avg:34.89ms
step:497/1575 train_time:17336ms step_avg:34.88ms
step:498/1575 train_time:17375ms step_avg:34.89ms
step:499/1575 train_time:17406ms step_avg:34.88ms
step:500/1575 train_time:17444ms step_avg:34.89ms
step:500/1575 val_loss:4.2393 train_time:17493ms step_avg:34.99ms
step:501/1575 train_time:17512ms step_avg:34.95ms
step:502/1575 train_time:17531ms step_avg:34.92ms
step:503/1575 train_time:17549ms step_avg:34.89ms
step:504/1575 train_time:17588ms step_avg:34.90ms
step:505/1575 train_time:17621ms step_avg:34.89ms
step:506/1575 train_time:17660ms step_avg:34.90ms
step:507/1575 train_time:17692ms step_avg:34.90ms
step:508/1575 train_time:17731ms step_avg:34.90ms
step:509/1575 train_time:17763ms step_avg:34.90ms
step:510/1575 train_time:17802ms step_avg:34.91ms
step:511/1575 train_time:17833ms step_avg:34.90ms
step:512/1575 train_time:17871ms step_avg:34.90ms
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:18063ms step_avg:35.07ms
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:18243ms step_avg:35.22ms
step:519/1575 train_time:18307ms step_avg:35.27ms
step:520/1575 train_time:18364ms step_avg:35.32ms
step:521/1575 train_time:18427ms step_avg:35.37ms
step:522/1575 train_time:18486ms step_avg:35.41ms
step:523/1575 train_time:18550ms step_avg:35.47ms
step:524/1575 train_time:18610ms step_avg:35.51ms
step:525/1575 train_time:18674ms step_avg:35.57ms
step:526/1575 train_time:18734ms step_avg:35.62ms
step:527/1575 train_time:18799ms step_avg:35.67ms
step:528/1575 train_time:18858ms step_avg:35.72ms
step:529/1575 train_time:18922ms step_avg:35.77ms
step:530/1575 train_time:18983ms step_avg:35.82ms
step:531/1575 train_time:19044ms step_avg:35.86ms
step:532/1575 train_time:19103ms step_avg:35.91ms
step:533/1575 train_time:19166ms step_avg:35.96ms
step:534/1575 train_time:19225ms step_avg:36.00ms
step:535/1575 train_time:19288ms step_avg:36.05ms
step:536/1575 train_time:19347ms step_avg:36.10ms
step:537/1575 train_time:19410ms step_avg:36.14ms
step:538/1575 train_time:19469ms step_avg:36.19ms
step:539/1575 train_time:19532ms step_avg:36.24ms
step:540/1575 train_time:19592ms step_avg:36.28ms
step:541/1575 train_time:19656ms step_avg:36.33ms
step:542/1575 train_time:19716ms step_avg:36.38ms
step:543/1575 train_time:19781ms step_avg:36.43ms
step:544/1575 train_time:19840ms step_avg:36.47ms
step:545/1575 train_time:19904ms step_avg:36.52ms
step:546/1575 train_time:19964ms step_avg:36.56ms
step:547/1575 train_time:20027ms step_avg:36.61ms
step:548/1575 train_time:20086ms step_avg:36.65ms
step:549/1575 train_time:20149ms step_avg:36.70ms
step:550/1575 train_time:20208ms step_avg:36.74ms
step:551/1575 train_time:20270ms step_avg:36.79ms
step:552/1575 train_time:20329ms step_avg:36.83ms
step:553/1575 train_time:20392ms step_avg:36.88ms
step:554/1575 train_time:20451ms step_avg:36.92ms
step:555/1575 train_time:20514ms step_avg:36.96ms
step:556/1575 train_time:20577ms step_avg:37.01ms
step:557/1575 train_time:20640ms step_avg:37.06ms
step:558/1575 train_time:20697ms step_avg:37.09ms
step:559/1575 train_time:20761ms step_avg:37.14ms
step:560/1575 train_time:20820ms step_avg:37.18ms
step:561/1575 train_time:20884ms step_avg:37.23ms
step:562/1575 train_time:20943ms step_avg:37.27ms
step:563/1575 train_time:21006ms step_avg:37.31ms
step:564/1575 train_time:21065ms step_avg:37.35ms
step:565/1575 train_time:21130ms step_avg:37.40ms
step:566/1575 train_time:21188ms step_avg:37.43ms
step:567/1575 train_time:21250ms step_avg:37.48ms
step:568/1575 train_time:21309ms step_avg:37.52ms
step:569/1575 train_time:21377ms step_avg:37.57ms
step:570/1575 train_time:21433ms step_avg:37.60ms
step:571/1575 train_time:21495ms step_avg:37.64ms
step:572/1575 train_time:21555ms step_avg:37.68ms
step:573/1575 train_time:21617ms step_avg:37.73ms
step:574/1575 train_time:21677ms step_avg:37.76ms
step:575/1575 train_time:21741ms step_avg:37.81ms
step:576/1575 train_time:21800ms step_avg:37.85ms
step:577/1575 train_time:21864ms step_avg:37.89ms
step:578/1575 train_time:21923ms step_avg:37.93ms
step:579/1575 train_time:21986ms step_avg:37.97ms
step:580/1575 train_time:22045ms step_avg:38.01ms
step:581/1575 train_time:22109ms step_avg:38.05ms
step:582/1575 train_time:22168ms step_avg:38.09ms
step:583/1575 train_time:22231ms step_avg:38.13ms
step:584/1575 train_time:22290ms step_avg:38.17ms
step:585/1575 train_time:22354ms step_avg:38.21ms
step:586/1575 train_time:22413ms step_avg:38.25ms
step:587/1575 train_time:22476ms step_avg:38.29ms
step:588/1575 train_time:22535ms step_avg:38.32ms
step:589/1575 train_time:22598ms step_avg:38.37ms
step:590/1575 train_time:22657ms step_avg:38.40ms
step:591/1575 train_time:22720ms step_avg:38.44ms
step:592/1575 train_time:22780ms step_avg:38.48ms
step:593/1575 train_time:22844ms step_avg:38.52ms
step:594/1575 train_time:22905ms step_avg:38.56ms
step:595/1575 train_time:22968ms step_avg:38.60ms
step:596/1575 train_time:23028ms step_avg:38.64ms
step:597/1575 train_time:23092ms step_avg:38.68ms
step:598/1575 train_time:23151ms step_avg:38.71ms
step:599/1575 train_time:23213ms step_avg:38.75ms
step:600/1575 train_time:23272ms step_avg:38.79ms
step:601/1575 train_time:23335ms step_avg:38.83ms
step:602/1575 train_time:23395ms step_avg:38.86ms
step:603/1575 train_time:23457ms step_avg:38.90ms
step:604/1575 train_time:23517ms step_avg:38.94ms
step:605/1575 train_time:23580ms step_avg:38.98ms
step:606/1575 train_time:23642ms step_avg:39.01ms
step:607/1575 train_time:23703ms step_avg:39.05ms
step:608/1575 train_time:23762ms step_avg:39.08ms
step:609/1575 train_time:23825ms step_avg:39.12ms
step:610/1575 train_time:23885ms step_avg:39.16ms
step:611/1575 train_time:23948ms step_avg:39.19ms
step:612/1575 train_time:24007ms step_avg:39.23ms
step:613/1575 train_time:24071ms step_avg:39.27ms
step:614/1575 train_time:24132ms step_avg:39.30ms
step:615/1575 train_time:24194ms step_avg:39.34ms
step:616/1575 train_time:24252ms step_avg:39.37ms
step:617/1575 train_time:24315ms step_avg:39.41ms
step:618/1575 train_time:24374ms step_avg:39.44ms
step:619/1575 train_time:24437ms step_avg:39.48ms
step:620/1575 train_time:24496ms step_avg:39.51ms
step:621/1575 train_time:24560ms step_avg:39.55ms
step:622/1575 train_time:24619ms step_avg:39.58ms
step:623/1575 train_time:24683ms step_avg:39.62ms
step:624/1575 train_time:24742ms step_avg:39.65ms
step:625/1575 train_time:24806ms step_avg:39.69ms
step:626/1575 train_time:24865ms step_avg:39.72ms
step:627/1575 train_time:24928ms step_avg:39.76ms
step:628/1575 train_time:24988ms step_avg:39.79ms
step:629/1575 train_time:25051ms step_avg:39.83ms
step:630/1575 train_time:25111ms step_avg:39.86ms
step:631/1575 train_time:25174ms step_avg:39.90ms
step:632/1575 train_time:25233ms step_avg:39.93ms
step:633/1575 train_time:25296ms step_avg:39.96ms
step:634/1575 train_time:25355ms step_avg:39.99ms
step:635/1575 train_time:25419ms step_avg:40.03ms
step:636/1575 train_time:25478ms step_avg:40.06ms
step:637/1575 train_time:25542ms step_avg:40.10ms
step:638/1575 train_time:25601ms step_avg:40.13ms
step:639/1575 train_time:25666ms step_avg:40.17ms
step:640/1575 train_time:25727ms step_avg:40.20ms
step:641/1575 train_time:25787ms step_avg:40.23ms
step:642/1575 train_time:25846ms step_avg:40.26ms
step:643/1575 train_time:25909ms step_avg:40.29ms
step:644/1575 train_time:25968ms step_avg:40.32ms
step:645/1575 train_time:26032ms step_avg:40.36ms
step:646/1575 train_time:26091ms step_avg:40.39ms
step:647/1575 train_time:26154ms step_avg:40.42ms
step:648/1575 train_time:26213ms step_avg:40.45ms
step:649/1575 train_time:26276ms step_avg:40.49ms
step:650/1575 train_time:26335ms step_avg:40.52ms
step:651/1575 train_time:26398ms step_avg:40.55ms
step:652/1575 train_time:26457ms step_avg:40.58ms
step:653/1575 train_time:26521ms step_avg:40.61ms
step:654/1575 train_time:26581ms step_avg:40.64ms
step:655/1575 train_time:26646ms step_avg:40.68ms
step:656/1575 train_time:26704ms step_avg:40.71ms
step:657/1575 train_time:26767ms step_avg:40.74ms
step:658/1575 train_time:26826ms step_avg:40.77ms
step:659/1575 train_time:26890ms step_avg:40.80ms
step:660/1575 train_time:26949ms step_avg:40.83ms
step:661/1575 train_time:27012ms step_avg:40.87ms
step:662/1575 train_time:27071ms step_avg:40.89ms
step:663/1575 train_time:27136ms step_avg:40.93ms
step:664/1575 train_time:27194ms step_avg:40.95ms
step:665/1575 train_time:27257ms step_avg:40.99ms
step:666/1575 train_time:27316ms step_avg:41.02ms
step:667/1575 train_time:27380ms step_avg:41.05ms
step:668/1575 train_time:27439ms step_avg:41.08ms
step:669/1575 train_time:27502ms step_avg:41.11ms
step:670/1575 train_time:27562ms step_avg:41.14ms
step:671/1575 train_time:27625ms step_avg:41.17ms
step:672/1575 train_time:27684ms step_avg:41.20ms
step:673/1575 train_time:27749ms step_avg:41.23ms
step:674/1575 train_time:27810ms step_avg:41.26ms
step:675/1575 train_time:27871ms step_avg:41.29ms
step:676/1575 train_time:27930ms step_avg:41.32ms
step:677/1575 train_time:27994ms step_avg:41.35ms
step:678/1575 train_time:28053ms step_avg:41.38ms
step:679/1575 train_time:28116ms step_avg:41.41ms
step:680/1575 train_time:28175ms step_avg:41.43ms
step:681/1575 train_time:28238ms step_avg:41.47ms
step:682/1575 train_time:28298ms step_avg:41.49ms
step:683/1575 train_time:28361ms step_avg:41.52ms
step:684/1575 train_time:28421ms step_avg:41.55ms
step:685/1575 train_time:28484ms step_avg:41.58ms
step:686/1575 train_time:28543ms step_avg:41.61ms
step:687/1575 train_time:28607ms step_avg:41.64ms
step:688/1575 train_time:28666ms step_avg:41.67ms
step:689/1575 train_time:28729ms step_avg:41.70ms
step:690/1575 train_time:28789ms step_avg:41.72ms
step:691/1575 train_time:28852ms step_avg:41.75ms
step:692/1575 train_time:28911ms step_avg:41.78ms
step:693/1575 train_time:28974ms step_avg:41.81ms
step:694/1575 train_time:29034ms step_avg:41.84ms
step:695/1575 train_time:29097ms step_avg:41.87ms
step:696/1575 train_time:29157ms step_avg:41.89ms
step:697/1575 train_time:29219ms step_avg:41.92ms
step:698/1575 train_time:29279ms step_avg:41.95ms
step:699/1575 train_time:29342ms step_avg:41.98ms
step:700/1575 train_time:29402ms step_avg:42.00ms
step:701/1575 train_time:29464ms step_avg:42.03ms
step:702/1575 train_time:29524ms step_avg:42.06ms
step:703/1575 train_time:29587ms step_avg:42.09ms
step:704/1575 train_time:29648ms step_avg:42.11ms
step:705/1575 train_time:29709ms step_avg:42.14ms
step:706/1575 train_time:29768ms step_avg:42.16ms
step:707/1575 train_time:29832ms step_avg:42.20ms
step:708/1575 train_time:29892ms step_avg:42.22ms
step:709/1575 train_time:29955ms step_avg:42.25ms
step:710/1575 train_time:30015ms step_avg:42.27ms
step:711/1575 train_time:30077ms step_avg:42.30ms
step:712/1575 train_time:30138ms step_avg:42.33ms
step:713/1575 train_time:30201ms step_avg:42.36ms
step:714/1575 train_time:30259ms step_avg:42.38ms
step:715/1575 train_time:30322ms step_avg:42.41ms
step:716/1575 train_time:30382ms step_avg:42.43ms
step:717/1575 train_time:30445ms step_avg:42.46ms
step:718/1575 train_time:30505ms step_avg:42.49ms
step:719/1575 train_time:30568ms step_avg:42.52ms
step:720/1575 train_time:30628ms step_avg:42.54ms
step:721/1575 train_time:30692ms step_avg:42.57ms
step:722/1575 train_time:30751ms step_avg:42.59ms
step:723/1575 train_time:30815ms step_avg:42.62ms
step:724/1575 train_time:30874ms step_avg:42.64ms
step:725/1575 train_time:30937ms step_avg:42.67ms
step:726/1575 train_time:30996ms step_avg:42.69ms
step:727/1575 train_time:31059ms step_avg:42.72ms
step:728/1575 train_time:31118ms step_avg:42.74ms
step:729/1575 train_time:31182ms step_avg:42.77ms
step:730/1575 train_time:31242ms step_avg:42.80ms
step:731/1575 train_time:31305ms step_avg:42.82ms
step:732/1575 train_time:31363ms step_avg:42.85ms
step:733/1575 train_time:31427ms step_avg:42.87ms
step:734/1575 train_time:31486ms step_avg:42.90ms
step:735/1575 train_time:31550ms step_avg:42.92ms
step:736/1575 train_time:31609ms step_avg:42.95ms
step:737/1575 train_time:31672ms step_avg:42.97ms
step:738/1575 train_time:31731ms step_avg:43.00ms
step:739/1575 train_time:31795ms step_avg:43.02ms
step:740/1575 train_time:31855ms step_avg:43.05ms
step:741/1575 train_time:31918ms step_avg:43.07ms
step:742/1575 train_time:31977ms step_avg:43.10ms
step:743/1575 train_time:32042ms step_avg:43.12ms
step:744/1575 train_time:32101ms step_avg:43.15ms
step:745/1575 train_time:32164ms step_avg:43.17ms
step:746/1575 train_time:32224ms step_avg:43.20ms
step:747/1575 train_time:32286ms step_avg:43.22ms
step:748/1575 train_time:32346ms step_avg:43.24ms
step:749/1575 train_time:32409ms step_avg:43.27ms
step:750/1575 train_time:32467ms step_avg:43.29ms
step:750/1575 val_loss:3.8754 train_time:32515ms step_avg:43.35ms
step:751/1575 train_time:32534ms step_avg:43.32ms
step:752/1575 train_time:32594ms step_avg:43.34ms
step:753/1575 train_time:32658ms step_avg:43.37ms
step:754/1575 train_time:32721ms step_avg:43.40ms
step:755/1575 train_time:32785ms step_avg:43.42ms
step:756/1575 train_time:32844ms step_avg:43.44ms
step:757/1575 train_time:32906ms step_avg:43.47ms
step:758/1575 train_time:32965ms step_avg:43.49ms
step:759/1575 train_time:33028ms step_avg:43.52ms
step:760/1575 train_time:33087ms step_avg:43.54ms
step:761/1575 train_time:33150ms step_avg:43.56ms
step:762/1575 train_time:33209ms step_avg:43.58ms
step:763/1575 train_time:33273ms step_avg:43.61ms
step:764/1575 train_time:33331ms step_avg:43.63ms
step:765/1575 train_time:33394ms step_avg:43.65ms
step:766/1575 train_time:33453ms step_avg:43.67ms
step:767/1575 train_time:33517ms step_avg:43.70ms
step:768/1575 train_time:33578ms step_avg:43.72ms
step:769/1575 train_time:33642ms step_avg:43.75ms
step:770/1575 train_time:33703ms step_avg:43.77ms
step:771/1575 train_time:33767ms step_avg:43.80ms
step:772/1575 train_time:33827ms step_avg:43.82ms
step:773/1575 train_time:33891ms step_avg:43.84ms
step:774/1575 train_time:33950ms step_avg:43.86ms
step:775/1575 train_time:34014ms step_avg:43.89ms
step:776/1575 train_time:34074ms step_avg:43.91ms
step:777/1575 train_time:34137ms step_avg:43.93ms
step:778/1575 train_time:34196ms step_avg:43.95ms
step:779/1575 train_time:34258ms step_avg:43.98ms
step:780/1575 train_time:34317ms step_avg:44.00ms
step:781/1575 train_time:34380ms step_avg:44.02ms
step:782/1575 train_time:34440ms step_avg:44.04ms
step:783/1575 train_time:34502ms step_avg:44.06ms
step:784/1575 train_time:34562ms step_avg:44.08ms
step:785/1575 train_time:34626ms step_avg:44.11ms
step:786/1575 train_time:34686ms step_avg:44.13ms
step:787/1575 train_time:34749ms step_avg:44.15ms
step:788/1575 train_time:34808ms step_avg:44.17ms
step:789/1575 train_time:34872ms step_avg:44.20ms
step:790/1575 train_time:34933ms step_avg:44.22ms
step:791/1575 train_time:34996ms step_avg:44.24ms
step:792/1575 train_time:35057ms step_avg:44.26ms
step:793/1575 train_time:35120ms step_avg:44.29ms
step:794/1575 train_time:35177ms step_avg:44.30ms
step:795/1575 train_time:35240ms step_avg:44.33ms
step:796/1575 train_time:35299ms step_avg:44.35ms
step:797/1575 train_time:35362ms step_avg:44.37ms
step:798/1575 train_time:35421ms step_avg:44.39ms
step:799/1575 train_time:35484ms step_avg:44.41ms
step:800/1575 train_time:35543ms step_avg:44.43ms
step:801/1575 train_time:35606ms step_avg:44.45ms
step:802/1575 train_time:35666ms step_avg:44.47ms
step:803/1575 train_time:35730ms step_avg:44.50ms
step:804/1575 train_time:35790ms step_avg:44.51ms
step:805/1575 train_time:35852ms step_avg:44.54ms
step:806/1575 train_time:35912ms step_avg:44.56ms
step:807/1575 train_time:35976ms step_avg:44.58ms
step:808/1575 train_time:36035ms step_avg:44.60ms
step:809/1575 train_time:36099ms step_avg:44.62ms
step:810/1575 train_time:36158ms step_avg:44.64ms
step:811/1575 train_time:36221ms step_avg:44.66ms
step:812/1575 train_time:36282ms step_avg:44.68ms
step:813/1575 train_time:36343ms step_avg:44.70ms
step:814/1575 train_time:36402ms step_avg:44.72ms
step:815/1575 train_time:36465ms step_avg:44.74ms
step:816/1575 train_time:36525ms step_avg:44.76ms
step:817/1575 train_time:36588ms step_avg:44.78ms
step:818/1575 train_time:36647ms step_avg:44.80ms
step:819/1575 train_time:36710ms step_avg:44.82ms
step:820/1575 train_time:36770ms step_avg:44.84ms
step:821/1575 train_time:36833ms step_avg:44.86ms
step:822/1575 train_time:36893ms step_avg:44.88ms
step:823/1575 train_time:36956ms step_avg:44.90ms
step:824/1575 train_time:37017ms step_avg:44.92ms
step:825/1575 train_time:37080ms step_avg:44.95ms
step:826/1575 train_time:37139ms step_avg:44.96ms
step:827/1575 train_time:37202ms step_avg:44.98ms
step:828/1575 train_time:37262ms step_avg:45.00ms
step:829/1575 train_time:37325ms step_avg:45.02ms
step:830/1575 train_time:37384ms step_avg:45.04ms
step:831/1575 train_time:37447ms step_avg:45.06ms
step:832/1575 train_time:37507ms step_avg:45.08ms
step:833/1575 train_time:37569ms step_avg:45.10ms
step:834/1575 train_time:37628ms step_avg:45.12ms
step:835/1575 train_time:37691ms step_avg:45.14ms
step:836/1575 train_time:37751ms step_avg:45.16ms
step:837/1575 train_time:37814ms step_avg:45.18ms
step:838/1575 train_time:37874ms step_avg:45.20ms
step:839/1575 train_time:37938ms step_avg:45.22ms
step:840/1575 train_time:37998ms step_avg:45.24ms
step:841/1575 train_time:38061ms step_avg:45.26ms
step:842/1575 train_time:38120ms step_avg:45.27ms
step:843/1575 train_time:38183ms step_avg:45.29ms
step:844/1575 train_time:38242ms step_avg:45.31ms
step:845/1575 train_time:38306ms step_avg:45.33ms
step:846/1575 train_time:38365ms step_avg:45.35ms
step:847/1575 train_time:38428ms step_avg:45.37ms
step:848/1575 train_time:38487ms step_avg:45.39ms
step:849/1575 train_time:38550ms step_avg:45.41ms
step:850/1575 train_time:38609ms step_avg:45.42ms
step:851/1575 train_time:38673ms step_avg:45.44ms
step:852/1575 train_time:38732ms step_avg:45.46ms
step:853/1575 train_time:38796ms step_avg:45.48ms
step:854/1575 train_time:38855ms step_avg:45.50ms
step:855/1575 train_time:38918ms step_avg:45.52ms
step:856/1575 train_time:38977ms step_avg:45.53ms
step:857/1575 train_time:39041ms step_avg:45.56ms
step:858/1575 train_time:39101ms step_avg:45.57ms
step:859/1575 train_time:39164ms step_avg:45.59ms
step:860/1575 train_time:39224ms step_avg:45.61ms
step:861/1575 train_time:39288ms step_avg:45.63ms
step:862/1575 train_time:39346ms step_avg:45.65ms
step:863/1575 train_time:39409ms step_avg:45.67ms
step:864/1575 train_time:39469ms step_avg:45.68ms
step:865/1575 train_time:39532ms step_avg:45.70ms
step:866/1575 train_time:39591ms step_avg:45.72ms
step:867/1575 train_time:39654ms step_avg:45.74ms
step:868/1575 train_time:39714ms step_avg:45.75ms
step:869/1575 train_time:39778ms step_avg:45.77ms
step:870/1575 train_time:39837ms step_avg:45.79ms
step:871/1575 train_time:39900ms step_avg:45.81ms
step:872/1575 train_time:39959ms step_avg:45.82ms
step:873/1575 train_time:40022ms step_avg:45.84ms
step:874/1575 train_time:40081ms step_avg:45.86ms
step:875/1575 train_time:40146ms step_avg:45.88ms
step:876/1575 train_time:40204ms step_avg:45.89ms
step:877/1575 train_time:40267ms step_avg:45.91ms
step:878/1575 train_time:40326ms step_avg:45.93ms
step:879/1575 train_time:40389ms step_avg:45.95ms
step:880/1575 train_time:40449ms step_avg:45.96ms
step:881/1575 train_time:40512ms step_avg:45.98ms
step:882/1575 train_time:40571ms step_avg:46.00ms
step:883/1575 train_time:40634ms step_avg:46.02ms
step:884/1575 train_time:40693ms step_avg:46.03ms
step:885/1575 train_time:40757ms step_avg:46.05ms
step:886/1575 train_time:40821ms step_avg:46.07ms
step:887/1575 train_time:40883ms step_avg:46.09ms
step:888/1575 train_time:40942ms step_avg:46.11ms
step:889/1575 train_time:41005ms step_avg:46.13ms
step:890/1575 train_time:41065ms step_avg:46.14ms
step:891/1575 train_time:41128ms step_avg:46.16ms
step:892/1575 train_time:41187ms step_avg:46.17ms
step:893/1575 train_time:41251ms step_avg:46.19ms
step:894/1575 train_time:41310ms step_avg:46.21ms
step:895/1575 train_time:41373ms step_avg:46.23ms
step:896/1575 train_time:41432ms step_avg:46.24ms
step:897/1575 train_time:41495ms step_avg:46.26ms
step:898/1575 train_time:41554ms step_avg:46.27ms
step:899/1575 train_time:41617ms step_avg:46.29ms
step:900/1575 train_time:41678ms step_avg:46.31ms
step:901/1575 train_time:41741ms step_avg:46.33ms
step:902/1575 train_time:41801ms step_avg:46.34ms
step:903/1575 train_time:41862ms step_avg:46.36ms
step:904/1575 train_time:41921ms step_avg:46.37ms
step:905/1575 train_time:41985ms step_avg:46.39ms
step:906/1575 train_time:42044ms step_avg:46.41ms
step:907/1575 train_time:42107ms step_avg:46.42ms
step:908/1575 train_time:42166ms step_avg:46.44ms
step:909/1575 train_time:42230ms step_avg:46.46ms
step:910/1575 train_time:42291ms step_avg:46.47ms
step:911/1575 train_time:42352ms step_avg:46.49ms
step:912/1575 train_time:42412ms step_avg:46.50ms
step:913/1575 train_time:42475ms step_avg:46.52ms
step:914/1575 train_time:42535ms step_avg:46.54ms
step:915/1575 train_time:42599ms step_avg:46.56ms
step:916/1575 train_time:42658ms step_avg:46.57ms
step:917/1575 train_time:42721ms step_avg:46.59ms
step:918/1575 train_time:42782ms step_avg:46.60ms
step:919/1575 train_time:42844ms step_avg:46.62ms
step:920/1575 train_time:42903ms step_avg:46.63ms
step:921/1575 train_time:42966ms step_avg:46.65ms
step:922/1575 train_time:43026ms step_avg:46.67ms
step:923/1575 train_time:43089ms step_avg:46.68ms
step:924/1575 train_time:43149ms step_avg:46.70ms
step:925/1575 train_time:43213ms step_avg:46.72ms
step:926/1575 train_time:43272ms step_avg:46.73ms
step:927/1575 train_time:43335ms step_avg:46.75ms
step:928/1575 train_time:43394ms step_avg:46.76ms
step:929/1575 train_time:43458ms step_avg:46.78ms
step:930/1575 train_time:43517ms step_avg:46.79ms
step:931/1575 train_time:43580ms step_avg:46.81ms
step:932/1575 train_time:43639ms step_avg:46.82ms
step:933/1575 train_time:43703ms step_avg:46.84ms
step:934/1575 train_time:43762ms step_avg:46.85ms
step:935/1575 train_time:43826ms step_avg:46.87ms
step:936/1575 train_time:43885ms step_avg:46.89ms
step:937/1575 train_time:43948ms step_avg:46.90ms
step:938/1575 train_time:44008ms step_avg:46.92ms
step:939/1575 train_time:44071ms step_avg:46.93ms
step:940/1575 train_time:44131ms step_avg:46.95ms
step:941/1575 train_time:44195ms step_avg:46.97ms
step:942/1575 train_time:44254ms step_avg:46.98ms
step:943/1575 train_time:44319ms step_avg:47.00ms
step:944/1575 train_time:44376ms step_avg:47.01ms
step:945/1575 train_time:44439ms step_avg:47.03ms
step:946/1575 train_time:44498ms step_avg:47.04ms
step:947/1575 train_time:44562ms step_avg:47.06ms
step:948/1575 train_time:44621ms step_avg:47.07ms
step:949/1575 train_time:44685ms step_avg:47.09ms
step:950/1575 train_time:44744ms step_avg:47.10ms
step:951/1575 train_time:44807ms step_avg:47.12ms
step:952/1575 train_time:44866ms step_avg:47.13ms
step:953/1575 train_time:44929ms step_avg:47.14ms
step:954/1575 train_time:44989ms step_avg:47.16ms
step:955/1575 train_time:45052ms step_avg:47.18ms
step:956/1575 train_time:45111ms step_avg:47.19ms
step:957/1575 train_time:45175ms step_avg:47.20ms
step:958/1575 train_time:45234ms step_avg:47.22ms
step:959/1575 train_time:45299ms step_avg:47.24ms
step:960/1575 train_time:45357ms step_avg:47.25ms
step:961/1575 train_time:45421ms step_avg:47.26ms
step:962/1575 train_time:45480ms step_avg:47.28ms
step:963/1575 train_time:45543ms step_avg:47.29ms
step:964/1575 train_time:45602ms step_avg:47.31ms
step:965/1575 train_time:45666ms step_avg:47.32ms
step:966/1575 train_time:45726ms step_avg:47.33ms
step:967/1575 train_time:45790ms step_avg:47.35ms
step:968/1575 train_time:45848ms step_avg:47.36ms
step:969/1575 train_time:45911ms step_avg:47.38ms
step:970/1575 train_time:45970ms step_avg:47.39ms
step:971/1575 train_time:46034ms step_avg:47.41ms
step:972/1575 train_time:46093ms step_avg:47.42ms
step:973/1575 train_time:46156ms step_avg:47.44ms
step:974/1575 train_time:46216ms step_avg:47.45ms
step:975/1575 train_time:46280ms step_avg:47.47ms
step:976/1575 train_time:46340ms step_avg:47.48ms
step:977/1575 train_time:46402ms step_avg:47.49ms
step:978/1575 train_time:46461ms step_avg:47.51ms
step:979/1575 train_time:46524ms step_avg:47.52ms
step:980/1575 train_time:46584ms step_avg:47.53ms
step:981/1575 train_time:46647ms step_avg:47.55ms
step:982/1575 train_time:46706ms step_avg:47.56ms
step:983/1575 train_time:46769ms step_avg:47.58ms
step:984/1575 train_time:46829ms step_avg:47.59ms
step:985/1575 train_time:46894ms step_avg:47.61ms
step:986/1575 train_time:46952ms step_avg:47.62ms
step:987/1575 train_time:47014ms step_avg:47.63ms
step:988/1575 train_time:47074ms step_avg:47.65ms
step:989/1575 train_time:47137ms step_avg:47.66ms
step:990/1575 train_time:47197ms step_avg:47.67ms
step:991/1575 train_time:47260ms step_avg:47.69ms
step:992/1575 train_time:47319ms step_avg:47.70ms
step:993/1575 train_time:47383ms step_avg:47.72ms
step:994/1575 train_time:47444ms step_avg:47.73ms
step:995/1575 train_time:47505ms step_avg:47.74ms
step:996/1575 train_time:47565ms step_avg:47.76ms
step:997/1575 train_time:47628ms step_avg:47.77ms
step:998/1575 train_time:47687ms step_avg:47.78ms
step:999/1575 train_time:47750ms step_avg:47.80ms
step:1000/1575 train_time:47810ms step_avg:47.81ms
step:1000/1575 val_loss:3.5853 train_time:47857ms step_avg:47.86ms
step:1001/1575 train_time:47877ms step_avg:47.83ms
step:1002/1575 train_time:47936ms step_avg:47.84ms
step:1003/1575 train_time:48000ms step_avg:47.86ms
step:1004/1575 train_time:48062ms step_avg:47.87ms
step:1005/1575 train_time:48126ms step_avg:47.89ms
step:1006/1575 train_time:48185ms step_avg:47.90ms
step:1007/1575 train_time:48248ms step_avg:47.91ms
step:1008/1575 train_time:48308ms step_avg:47.92ms
step:1009/1575 train_time:48370ms step_avg:47.94ms
step:1010/1575 train_time:48431ms step_avg:47.95ms
step:1011/1575 train_time:48493ms step_avg:47.97ms
step:1012/1575 train_time:48552ms step_avg:47.98ms
step:1013/1575 train_time:48614ms step_avg:47.99ms
step:1014/1575 train_time:48673ms step_avg:48.00ms
step:1015/1575 train_time:48736ms step_avg:48.02ms
step:1016/1575 train_time:48795ms step_avg:48.03ms
step:1017/1575 train_time:48859ms step_avg:48.04ms
step:1018/1575 train_time:48920ms step_avg:48.05ms
step:1019/1575 train_time:48983ms step_avg:48.07ms
step:1020/1575 train_time:49043ms step_avg:48.08ms
step:1021/1575 train_time:49109ms step_avg:48.10ms
step:1022/1575 train_time:49169ms step_avg:48.11ms
step:1023/1575 train_time:49230ms step_avg:48.12ms
step:1024/1575 train_time:49291ms step_avg:48.14ms
step:1025/1575 train_time:49361ms step_avg:48.16ms
step:1026/1575 train_time:49445ms step_avg:48.19ms
step:1027/1575 train_time:49535ms step_avg:48.23ms
step:1028/1575 train_time:49620ms step_avg:48.27ms
step:1029/1575 train_time:49708ms step_avg:48.31ms
step:1030/1575 train_time:49794ms step_avg:48.34ms
step:1031/1575 train_time:49885ms step_avg:48.39ms
step:1032/1575 train_time:49971ms step_avg:48.42ms
step:1033/1575 train_time:50061ms step_avg:48.46ms
step:1034/1575 train_time:50147ms step_avg:48.50ms
step:1035/1575 train_time:50239ms step_avg:48.54ms
step:1036/1575 train_time:50325ms step_avg:48.58ms
step:1037/1575 train_time:50415ms step_avg:48.62ms
step:1038/1575 train_time:50500ms step_avg:48.65ms
step:1039/1575 train_time:50590ms step_avg:48.69ms
step:1040/1575 train_time:50675ms step_avg:48.73ms
step:1041/1575 train_time:50765ms step_avg:48.77ms
step:1042/1575 train_time:50851ms step_avg:48.80ms
step:1043/1575 train_time:50941ms step_avg:48.84ms
step:1044/1575 train_time:51027ms step_avg:48.88ms
step:1045/1575 train_time:51118ms step_avg:48.92ms
step:1046/1575 train_time:51203ms step_avg:48.95ms
step:1047/1575 train_time:51293ms step_avg:48.99ms
step:1048/1575 train_time:51379ms step_avg:49.03ms
step:1049/1575 train_time:51468ms step_avg:49.06ms
step:1050/1575 train_time:51555ms step_avg:49.10ms
step:1051/1575 train_time:51642ms step_avg:49.14ms
step:1052/1575 train_time:51727ms step_avg:49.17ms
step:1053/1575 train_time:51818ms step_avg:49.21ms
step:1054/1575 train_time:51904ms step_avg:49.24ms
step:1055/1575 train_time:51994ms step_avg:49.28ms
step:1056/1575 train_time:52079ms step_avg:49.32ms
step:1057/1575 train_time:52169ms step_avg:49.36ms
step:1058/1575 train_time:52256ms step_avg:49.39ms
step:1059/1575 train_time:52346ms step_avg:49.43ms
step:1060/1575 train_time:52434ms step_avg:49.47ms
step:1061/1575 train_time:52522ms step_avg:49.50ms
step:1062/1575 train_time:52608ms step_avg:49.54ms
step:1063/1575 train_time:52697ms step_avg:49.57ms
step:1064/1575 train_time:52783ms step_avg:49.61ms
step:1065/1575 train_time:52873ms step_avg:49.65ms
step:1066/1575 train_time:52960ms step_avg:49.68ms
step:1067/1575 train_time:53049ms step_avg:49.72ms
step:1068/1575 train_time:53135ms step_avg:49.75ms
step:1069/1575 train_time:53225ms step_avg:49.79ms
step:1070/1575 train_time:53310ms step_avg:49.82ms
step:1071/1575 train_time:53400ms step_avg:49.86ms
step:1072/1575 train_time:53485ms step_avg:49.89ms
step:1073/1575 train_time:53576ms step_avg:49.93ms
step:1074/1575 train_time:53661ms step_avg:49.96ms
step:1075/1575 train_time:53750ms step_avg:50.00ms
step:1076/1575 train_time:53836ms step_avg:50.03ms
step:1077/1575 train_time:53926ms step_avg:50.07ms
step:1078/1575 train_time:54011ms step_avg:50.10ms
step:1079/1575 train_time:54101ms step_avg:50.14ms
step:1080/1575 train_time:54187ms step_avg:50.17ms
step:1081/1575 train_time:54277ms step_avg:50.21ms
step:1082/1575 train_time:54364ms step_avg:50.24ms
step:1083/1575 train_time:54453ms step_avg:50.28ms
step:1084/1575 train_time:54539ms step_avg:50.31ms
step:1085/1575 train_time:54635ms step_avg:50.35ms
step:1086/1575 train_time:54722ms step_avg:50.39ms
step:1087/1575 train_time:54805ms step_avg:50.42ms
step:1088/1575 train_time:54890ms step_avg:50.45ms
step:1089/1575 train_time:54980ms step_avg:50.49ms
step:1090/1575 train_time:55066ms step_avg:50.52ms
step:1091/1575 train_time:55156ms step_avg:50.56ms
step:1092/1575 train_time:55241ms step_avg:50.59ms
step:1093/1575 train_time:55330ms step_avg:50.62ms
step:1094/1575 train_time:55416ms step_avg:50.65ms
step:1095/1575 train_time:55505ms step_avg:50.69ms
step:1096/1575 train_time:55592ms step_avg:50.72ms
step:1097/1575 train_time:55682ms step_avg:50.76ms
step:1098/1575 train_time:55767ms step_avg:50.79ms
step:1099/1575 train_time:55857ms step_avg:50.83ms
step:1100/1575 train_time:55944ms step_avg:50.86ms
step:1101/1575 train_time:56034ms step_avg:50.89ms
step:1102/1575 train_time:56119ms step_avg:50.92ms
step:1103/1575 train_time:56218ms step_avg:50.97ms
step:1104/1575 train_time:56295ms step_avg:50.99ms
step:1105/1575 train_time:56384ms step_avg:51.03ms
step:1106/1575 train_time:56470ms step_avg:51.06ms
step:1107/1575 train_time:56560ms step_avg:51.09ms
step:1108/1575 train_time:56645ms step_avg:51.12ms
step:1109/1575 train_time:56735ms step_avg:51.16ms
step:1110/1575 train_time:56824ms step_avg:51.19ms
step:1111/1575 train_time:56911ms step_avg:51.23ms
step:1112/1575 train_time:57000ms step_avg:51.26ms
step:1113/1575 train_time:57086ms step_avg:51.29ms
step:1114/1575 train_time:57172ms step_avg:51.32ms
step:1115/1575 train_time:57261ms step_avg:51.36ms
step:1116/1575 train_time:57347ms step_avg:51.39ms
step:1117/1575 train_time:57439ms step_avg:51.42ms
step:1118/1575 train_time:57528ms step_avg:51.46ms
step:1119/1575 train_time:57614ms step_avg:51.49ms
step:1120/1575 train_time:57702ms step_avg:51.52ms
step:1121/1575 train_time:57789ms step_avg:51.55ms
step:1122/1575 train_time:57874ms step_avg:51.58ms
step:1123/1575 train_time:57963ms step_avg:51.61ms
step:1124/1575 train_time:58049ms step_avg:51.64ms
step:1125/1575 train_time:58139ms step_avg:51.68ms
step:1126/1575 train_time:58225ms step_avg:51.71ms
step:1127/1575 train_time:58315ms step_avg:51.74ms
step:1128/1575 train_time:58400ms step_avg:51.77ms
step:1129/1575 train_time:58489ms step_avg:51.81ms
step:1130/1575 train_time:58575ms step_avg:51.84ms
step:1131/1575 train_time:58665ms step_avg:51.87ms
step:1132/1575 train_time:58750ms step_avg:51.90ms
step:1133/1575 train_time:58843ms step_avg:51.94ms
step:1134/1575 train_time:58927ms step_avg:51.96ms
step:1135/1575 train_time:59015ms step_avg:52.00ms
step:1136/1575 train_time:59101ms step_avg:52.03ms
step:1137/1575 train_time:59191ms step_avg:52.06ms
step:1138/1575 train_time:59277ms step_avg:52.09ms
step:1139/1575 train_time:59365ms step_avg:52.12ms
step:1140/1575 train_time:59454ms step_avg:52.15ms
step:1141/1575 train_time:59543ms step_avg:52.18ms
step:1142/1575 train_time:59628ms step_avg:52.21ms
step:1143/1575 train_time:59717ms step_avg:52.25ms
step:1144/1575 train_time:59803ms step_avg:52.27ms
step:1145/1575 train_time:59892ms step_avg:52.31ms
step:1146/1575 train_time:59983ms step_avg:52.34ms
step:1147/1575 train_time:60072ms step_avg:52.37ms
step:1148/1575 train_time:60155ms step_avg:52.40ms
step:1149/1575 train_time:60245ms step_avg:52.43ms
step:1150/1575 train_time:60330ms step_avg:52.46ms
step:1151/1575 train_time:60420ms step_avg:52.49ms
step:1152/1575 train_time:60506ms step_avg:52.52ms
step:1153/1575 train_time:60595ms step_avg:52.55ms
step:1154/1575 train_time:60682ms step_avg:52.58ms
step:1155/1575 train_time:60771ms step_avg:52.62ms
step:1156/1575 train_time:60857ms step_avg:52.64ms
step:1157/1575 train_time:60948ms step_avg:52.68ms
step:1158/1575 train_time:61033ms step_avg:52.71ms
step:1159/1575 train_time:61122ms step_avg:52.74ms
step:1160/1575 train_time:61207ms step_avg:52.77ms
step:1161/1575 train_time:61297ms step_avg:52.80ms
step:1162/1575 train_time:61383ms step_avg:52.83ms
step:1163/1575 train_time:61472ms step_avg:52.86ms
step:1164/1575 train_time:61558ms step_avg:52.88ms
step:1165/1575 train_time:61647ms step_avg:52.92ms
step:1166/1575 train_time:61733ms step_avg:52.94ms
step:1167/1575 train_time:61824ms step_avg:52.98ms
step:1168/1575 train_time:61911ms step_avg:53.01ms
step:1169/1575 train_time:61999ms step_avg:53.04ms
step:1170/1575 train_time:62085ms step_avg:53.06ms
step:1171/1575 train_time:62174ms 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:62438ms step_avg:53.18ms
step:1175/1575 train_time:62526ms step_avg:53.21ms
step:1176/1575 train_time:62612ms step_avg:53.24ms
step:1177/1575 train_time:62701ms step_avg:53.27ms
step:1178/1575 train_time:62787ms step_avg:53.30ms
step:1179/1575 train_time:62877ms step_avg:53.33ms
step:1180/1575 train_time:62963ms step_avg:53.36ms
step:1181/1575 train_time:63052ms step_avg:53.39ms
step:1182/1575 train_time:63139ms step_avg:53.42ms
step:1183/1575 train_time:63228ms step_avg:53.45ms
step:1184/1575 train_time:63313ms step_avg:53.47ms
step:1185/1575 train_time:63404ms step_avg:53.51ms
step:1186/1575 train_time:63490ms step_avg:53.53ms
step:1187/1575 train_time:63580ms step_avg:53.56ms
step:1188/1575 train_time:63666ms step_avg:53.59ms
step:1189/1575 train_time:63756ms step_avg:53.62ms
step:1190/1575 train_time:63841ms step_avg:53.65ms
step:1191/1575 train_time:63931ms step_avg:53.68ms
step:1192/1575 train_time:64019ms step_avg:53.71ms
step:1193/1575 train_time:64106ms step_avg:53.74ms
step:1194/1575 train_time:64192ms step_avg:53.76ms
step:1195/1575 train_time:64282ms step_avg:53.79ms
step:1196/1575 train_time:64367ms step_avg:53.82ms
step:1197/1575 train_time:64457ms step_avg:53.85ms
step:1198/1575 train_time:64543ms step_avg:53.88ms
step:1199/1575 train_time:64632ms step_avg:53.90ms
step:1200/1575 train_time:64719ms step_avg:53.93ms
step:1201/1575 train_time:64808ms step_avg:53.96ms
step:1202/1575 train_time:64895ms step_avg:53.99ms
step:1203/1575 train_time:64984ms step_avg:54.02ms
step:1204/1575 train_time:65071ms step_avg:54.05ms
step:1205/1575 train_time:65160ms step_avg:54.07ms
step:1206/1575 train_time:65245ms step_avg:54.10ms
step:1207/1575 train_time:65335ms step_avg:54.13ms
step:1208/1575 train_time:65421ms step_avg:54.16ms
step:1209/1575 train_time:65513ms step_avg:54.19ms
step:1210/1575 train_time:65597ms step_avg:54.21ms
step:1211/1575 train_time:65687ms step_avg:54.24ms
step:1212/1575 train_time:65772ms step_avg:54.27ms
step:1213/1575 train_time:65862ms step_avg:54.30ms
step:1214/1575 train_time:65950ms step_avg:54.32ms
step:1215/1575 train_time:66037ms step_avg:54.35ms
step:1216/1575 train_time:66123ms step_avg:54.38ms
step:1217/1575 train_time:66211ms step_avg:54.41ms
step:1218/1575 train_time:66298ms step_avg:54.43ms
step:1219/1575 train_time:66387ms step_avg:54.46ms
step:1220/1575 train_time:66473ms step_avg:54.49ms
step:1221/1575 train_time:66563ms step_avg:54.51ms
step:1222/1575 train_time:66648ms step_avg:54.54ms
step:1223/1575 train_time:66738ms step_avg:54.57ms
step:1224/1575 train_time:66823ms step_avg:54.59ms
step:1225/1575 train_time:66913ms step_avg:54.62ms
step:1226/1575 train_time:66998ms step_avg:54.65ms
step:1227/1575 train_time:67088ms step_avg:54.68ms
step:1228/1575 train_time:67174ms step_avg:54.70ms
step:1229/1575 train_time:67264ms step_avg:54.73ms
step:1230/1575 train_time:67350ms step_avg:54.76ms
step:1231/1575 train_time:67440ms step_avg:54.79ms
step:1232/1575 train_time:67525ms step_avg:54.81ms
step:1233/1575 train_time:67615ms step_avg:54.84ms
step:1234/1575 train_time:67702ms step_avg:54.86ms
step:1235/1575 train_time:67790ms step_avg:54.89ms
step:1236/1575 train_time:67877ms step_avg:54.92ms
step:1237/1575 train_time:67967ms step_avg:54.94ms
step:1238/1575 train_time:68052ms step_avg:54.97ms
step:1239/1575 train_time:68142ms step_avg:55.00ms
step:1240/1575 train_time:68227ms step_avg:55.02ms
step:1241/1575 train_time:68317ms step_avg:55.05ms
step:1242/1575 train_time:68402ms step_avg:55.07ms
step:1243/1575 train_time:68492ms step_avg:55.10ms
step:1244/1575 train_time:68577ms step_avg:55.13ms
step:1245/1575 train_time:68666ms step_avg:55.15ms
step:1246/1575 train_time:68752ms step_avg:55.18ms
step:1247/1575 train_time:68842ms step_avg:55.21ms
step:1248/1575 train_time:68928ms step_avg:55.23ms
step:1249/1575 train_time:69018ms step_avg:55.26ms
step:1250/1575 train_time:69104ms step_avg:55.28ms
step:1250/1575 val_loss:3.4078 train_time:69177ms step_avg:55.34ms
step:1251/1575 train_time:69197ms step_avg:55.31ms
step:1252/1575 train_time:69285ms step_avg:55.34ms
step:1253/1575 train_time:69377ms step_avg:55.37ms
step:1254/1575 train_time:69463ms step_avg:55.39ms
step:1255/1575 train_time:69554ms step_avg:55.42ms
step:1256/1575 train_time:69638ms step_avg:55.44ms
step:1257/1575 train_time:69726ms step_avg:55.47ms
step:1258/1575 train_time:69811ms step_avg:55.49ms
step:1259/1575 train_time:69899ms step_avg:55.52ms
step:1260/1575 train_time:69984ms step_avg:55.54ms
step:1261/1575 train_time:70073ms step_avg:55.57ms
step:1262/1575 train_time:70158ms step_avg:55.59ms
step:1263/1575 train_time:70250ms step_avg:55.62ms
step:1264/1575 train_time:70337ms step_avg:55.65ms
step:1265/1575 train_time:70428ms step_avg:55.67ms
step:1266/1575 train_time:70514ms step_avg:55.70ms
step:1267/1575 train_time:70604ms step_avg:55.73ms
step:1268/1575 train_time:70689ms step_avg:55.75ms
step:1269/1575 train_time:70778ms step_avg:55.77ms
step:1270/1575 train_time:70862ms step_avg:55.80ms
step:1271/1575 train_time:70950ms step_avg:55.82ms
step:1272/1575 train_time:71035ms step_avg:55.85ms
step:1273/1575 train_time:71125ms step_avg:55.87ms
step:1274/1575 train_time:71210ms step_avg:55.90ms
step:1275/1575 train_time:71303ms step_avg:55.92ms
step:1276/1575 train_time:71390ms step_avg:55.95ms
step:1277/1575 train_time:71480ms step_avg:55.97ms
step:1278/1575 train_time:71565ms step_avg:56.00ms
step:1279/1575 train_time:71655ms step_avg:56.02ms
step:1280/1575 train_time:71740ms step_avg:56.05ms
step:1281/1575 train_time:71828ms step_avg:56.07ms
step:1282/1575 train_time:71913ms step_avg:56.09ms
step:1283/1575 train_time:72003ms step_avg:56.12ms
step:1284/1575 train_time:72088ms step_avg:56.14ms
step:1285/1575 train_time:72178ms step_avg:56.17ms
step:1286/1575 train_time:72266ms step_avg:56.19ms
step:1287/1575 train_time:72355ms step_avg:56.22ms
step:1288/1575 train_time:72443ms step_avg:56.24ms
step:1289/1575 train_time:72533ms step_avg:56.27ms
step:1290/1575 train_time:72619ms step_avg:56.29ms
step:1291/1575 train_time:72708ms step_avg:56.32ms
step:1292/1575 train_time:72794ms step_avg:56.34ms
step:1293/1575 train_time:72883ms step_avg:56.37ms
step:1294/1575 train_time:72968ms step_avg:56.39ms
step:1295/1575 train_time:73058ms step_avg:56.42ms
step:1296/1575 train_time:73143ms step_avg:56.44ms
step:1297/1575 train_time:73233ms step_avg:56.46ms
step:1298/1575 train_time:73319ms step_avg:56.49ms
step:1299/1575 train_time:73409ms step_avg:56.51ms
step:1300/1575 train_time:73495ms step_avg:56.53ms
step:1301/1575 train_time:73585ms step_avg:56.56ms
step:1302/1575 train_time:73670ms step_avg:56.58ms
step:1303/1575 train_time:73760ms step_avg:56.61ms
step:1304/1575 train_time:73844ms step_avg:56.63ms
step:1305/1575 train_time:73934ms step_avg:56.65ms
step:1306/1575 train_time:74019ms step_avg:56.68ms
step:1307/1575 train_time:74109ms step_avg:56.70ms
step:1308/1575 train_time:74195ms step_avg:56.72ms
step:1309/1575 train_time:74286ms step_avg:56.75ms
step:1310/1575 train_time:74371ms step_avg:56.77ms
step:1311/1575 train_time:74461ms step_avg:56.80ms
step:1312/1575 train_time:74547ms step_avg:56.82ms
step:1313/1575 train_time:74636ms step_avg:56.84ms
step:1314/1575 train_time:74722ms step_avg:56.87ms
step:1315/1575 train_time:74811ms step_avg:56.89ms
step:1316/1575 train_time:74898ms step_avg:56.91ms
step:1317/1575 train_time:74987ms step_avg:56.94ms
step:1318/1575 train_time:75072ms step_avg:56.96ms
step:1319/1575 train_time:75163ms step_avg:56.99ms
step:1320/1575 train_time:75250ms step_avg:57.01ms
step:1321/1575 train_time:75340ms step_avg:57.03ms
step:1322/1575 train_time:75426ms step_avg:57.05ms
step:1323/1575 train_time:75515ms step_avg:57.08ms
step:1324/1575 train_time:75601ms step_avg:57.10ms
step:1325/1575 train_time:75691ms step_avg:57.13ms
step:1326/1575 train_time:75779ms step_avg:57.15ms
step:1327/1575 train_time:75865ms step_avg:57.17ms
step:1328/1575 train_time:75951ms step_avg:57.19ms
step:1329/1575 train_time:76040ms step_avg:57.22ms
step:1330/1575 train_time:76125ms step_avg:57.24ms
step:1331/1575 train_time:76215ms step_avg:57.26ms
step:1332/1575 train_time:76301ms step_avg:57.28ms
step:1333/1575 train_time:76391ms step_avg:57.31ms
step:1334/1575 train_time:76477ms step_avg:57.33ms
step:1335/1575 train_time:76567ms step_avg:57.35ms
step:1336/1575 train_time:76652ms step_avg:57.37ms
step:1337/1575 train_time:76742ms step_avg:57.40ms
step:1338/1575 train_time:76828ms step_avg:57.42ms
step:1339/1575 train_time:76917ms step_avg:57.44ms
step:1340/1575 train_time:77003ms step_avg:57.46ms
step:1341/1575 train_time:77091ms step_avg:57.49ms
step:1342/1575 train_time:77178ms step_avg:57.51ms
step:1343/1575 train_time:77269ms step_avg:57.53ms
step:1344/1575 train_time:77354ms step_avg:57.55ms
step:1345/1575 train_time:77444ms step_avg:57.58ms
step:1346/1575 train_time:77530ms step_avg:57.60ms
step:1347/1575 train_time:77620ms step_avg:57.62ms
step:1348/1575 train_time:77705ms step_avg:57.65ms
step:1349/1575 train_time:77795ms step_avg:57.67ms
step:1350/1575 train_time:77881ms step_avg:57.69ms
step:1351/1575 train_time:77971ms step_avg:57.71ms
step:1352/1575 train_time:78055ms step_avg:57.73ms
step:1353/1575 train_time:78145ms step_avg:57.76ms
step:1354/1575 train_time:78232ms step_avg:57.78ms
step:1355/1575 train_time:78321ms step_avg:57.80ms
step:1356/1575 train_time:78406ms step_avg:57.82ms
step:1357/1575 train_time:78496ms step_avg:57.85ms
step:1358/1575 train_time:78586ms step_avg:57.87ms
step:1359/1575 train_time:78673ms step_avg:57.89ms
step:1360/1575 train_time:78760ms step_avg:57.91ms
step:1361/1575 train_time:78849ms step_avg:57.93ms
step:1362/1575 train_time:78933ms step_avg:57.95ms
step:1363/1575 train_time:79023ms step_avg:57.98ms
step:1364/1575 train_time:79108ms step_avg:58.00ms
step:1365/1575 train_time:79197ms step_avg:58.02ms
step:1366/1575 train_time:79285ms step_avg:58.04ms
step:1367/1575 train_time:79375ms step_avg:58.07ms
step:1368/1575 train_time:79460ms step_avg:58.09ms
step:1369/1575 train_time:79548ms step_avg:58.11ms
step:1370/1575 train_time:79635ms step_avg:58.13ms
step:1371/1575 train_time:79725ms step_avg:58.15ms
step:1372/1575 train_time:79811ms step_avg:58.17ms
step:1373/1575 train_time:79900ms step_avg:58.19ms
step:1374/1575 train_time:79986ms step_avg:58.21ms
step:1375/1575 train_time:80075ms step_avg:58.24ms
step:1376/1575 train_time:80160ms step_avg:58.26ms
step:1377/1575 train_time:80250ms step_avg:58.28ms
step:1378/1575 train_time:80336ms step_avg:58.30ms
step:1379/1575 train_time:80426ms step_avg:58.32ms
step:1380/1575 train_time:80512ms step_avg:58.34ms
step:1381/1575 train_time:80603ms step_avg:58.37ms
step:1382/1575 train_time:80690ms step_avg:58.39ms
step:1383/1575 train_time:80782ms step_avg:58.41ms
step:1384/1575 train_time:80867ms step_avg:58.43ms
step:1385/1575 train_time:80954ms step_avg:58.45ms
step:1386/1575 train_time:81041ms step_avg:58.47ms
step:1387/1575 train_time:81130ms step_avg:58.49ms
step:1388/1575 train_time:81215ms step_avg:58.51ms
step:1389/1575 train_time:81305ms step_avg:58.53ms
step:1390/1575 train_time:81391ms step_avg:58.55ms
step:1391/1575 train_time:81480ms step_avg:58.58ms
step:1392/1575 train_time:81565ms step_avg:58.60ms
step:1393/1575 train_time:81655ms step_avg:58.62ms
step:1394/1575 train_time:81741ms step_avg:58.64ms
step:1395/1575 train_time:81831ms step_avg:58.66ms
step:1396/1575 train_time:81917ms step_avg:58.68ms
step:1397/1575 train_time:82007ms step_avg:58.70ms
step:1398/1575 train_time:82092ms step_avg:58.72ms
step:1399/1575 train_time:82182ms step_avg:58.74ms
step:1400/1575 train_time:82267ms step_avg:58.76ms
step:1401/1575 train_time:82356ms step_avg:58.78ms
step:1402/1575 train_time:82442ms step_avg:58.80ms
step:1403/1575 train_time:82532ms step_avg:58.83ms
step:1404/1575 train_time:82618ms step_avg:58.84ms
step:1405/1575 train_time:82707ms step_avg:58.87ms
step:1406/1575 train_time:82794ms step_avg:58.89ms
step:1407/1575 train_time:82885ms step_avg:58.91ms
step:1408/1575 train_time:82969ms step_avg:58.93ms
step:1409/1575 train_time:83059ms step_avg:58.95ms
step:1410/1575 train_time:83145ms step_avg:58.97ms
step:1411/1575 train_time:83235ms step_avg:58.99ms
step:1412/1575 train_time:83321ms step_avg:59.01ms
step:1413/1575 train_time:83412ms step_avg:59.03ms
step:1414/1575 train_time:83498ms step_avg:59.05ms
step:1415/1575 train_time:83588ms step_avg:59.07ms
step:1416/1575 train_time:83673ms step_avg:59.09ms
step:1417/1575 train_time:83764ms step_avg:59.11ms
step:1418/1575 train_time:83849ms step_avg:59.13ms
step:1419/1575 train_time:83940ms step_avg:59.15ms
step:1420/1575 train_time:84024ms step_avg:59.17ms
step:1421/1575 train_time:84114ms step_avg:59.19ms
step:1422/1575 train_time:84200ms step_avg:59.21ms
step:1423/1575 train_time:84291ms step_avg:59.23ms
step:1424/1575 train_time:84375ms step_avg:59.25ms
step:1425/1575 train_time:84465ms step_avg:59.27ms
step:1426/1575 train_time:84554ms step_avg:59.29ms
step:1427/1575 train_time:84640ms step_avg:59.31ms
step:1428/1575 train_time:84726ms step_avg:59.33ms
step:1429/1575 train_time:84816ms step_avg:59.35ms
step:1430/1575 train_time:84902ms step_avg:59.37ms
step:1431/1575 train_time:84992ms step_avg:59.39ms
step:1432/1575 train_time:85077ms step_avg:59.41ms
step:1433/1575 train_time:85167ms step_avg:59.43ms
step:1434/1575 train_time:85253ms step_avg:59.45ms
step:1435/1575 train_time:85343ms step_avg:59.47ms
step:1436/1575 train_time:85428ms step_avg:59.49ms
step:1437/1575 train_time:85518ms step_avg:59.51ms
step:1438/1575 train_time:85603ms step_avg:59.53ms
step:1439/1575 train_time:85693ms step_avg:59.55ms
step:1440/1575 train_time:85780ms step_avg:59.57ms
step:1441/1575 train_time:85868ms step_avg:59.59ms
step:1442/1575 train_time:85954ms step_avg:59.61ms
step:1443/1575 train_time:86044ms step_avg:59.63ms
step:1444/1575 train_time:86130ms step_avg:59.65ms
step:1445/1575 train_time:86219ms step_avg:59.67ms
step:1446/1575 train_time:86306ms step_avg:59.69ms
step:1447/1575 train_time:86394ms step_avg:59.71ms
step:1448/1575 train_time:86480ms step_avg:59.72ms
step:1449/1575 train_time:86569ms step_avg:59.74ms
step:1450/1575 train_time:86656ms step_avg:59.76ms
step:1451/1575 train_time:86746ms step_avg:59.78ms
step:1452/1575 train_time:86831ms step_avg:59.80ms
step:1453/1575 train_time:86923ms step_avg:59.82ms
step:1454/1575 train_time:87007ms step_avg:59.84ms
step:1455/1575 train_time:87098ms step_avg:59.86ms
step:1456/1575 train_time:87184ms step_avg:59.88ms
step:1457/1575 train_time:87273ms step_avg:59.90ms
step:1458/1575 train_time:87359ms step_avg:59.92ms
step:1459/1575 train_time:87448ms step_avg:59.94ms
step:1460/1575 train_time:87535ms step_avg:59.96ms
step:1461/1575 train_time:87624ms step_avg:59.98ms
step:1462/1575 train_time:87710ms step_avg:59.99ms
step:1463/1575 train_time:87800ms step_avg:60.01ms
step:1464/1575 train_time:87885ms step_avg:60.03ms
step:1465/1575 train_time:87974ms step_avg:60.05ms
step:1466/1575 train_time:88060ms step_avg:60.07ms
step:1467/1575 train_time:88149ms step_avg:60.09ms
step:1468/1575 train_time:88236ms step_avg:60.11ms
step:1469/1575 train_time:88326ms step_avg:60.13ms
step:1470/1575 train_time:88412ms step_avg:60.14ms
step:1471/1575 train_time:88501ms step_avg:60.16ms
step:1472/1575 train_time:88586ms step_avg:60.18ms
step:1473/1575 train_time:88676ms step_avg:60.20ms
step:1474/1575 train_time:88761ms step_avg:60.22ms
step:1475/1575 train_time:88851ms step_avg:60.24ms
step:1476/1575 train_time:88938ms step_avg:60.26ms
step:1477/1575 train_time:89025ms step_avg:60.27ms
step:1478/1575 train_time:89111ms step_avg:60.29ms
step:1479/1575 train_time:89202ms step_avg:60.31ms
step:1480/1575 train_time:89289ms step_avg:60.33ms
step:1481/1575 train_time:89376ms step_avg:60.35ms
step:1482/1575 train_time:89461ms step_avg:60.37ms
step:1483/1575 train_time:89551ms step_avg:60.39ms
step:1484/1575 train_time:89641ms step_avg:60.41ms
step:1485/1575 train_time:89728ms step_avg:60.42ms
step:1486/1575 train_time:89814ms step_avg:60.44ms
step:1487/1575 train_time:89903ms step_avg:60.46ms
step:1488/1575 train_time:89989ms step_avg:60.48ms
step:1489/1575 train_time:90079ms step_avg:60.50ms
step:1490/1575 train_time:90164ms step_avg:60.51ms
step:1491/1575 train_time:90254ms step_avg:60.53ms
step:1492/1575 train_time:90340ms step_avg:60.55ms
step:1493/1575 train_time:90429ms step_avg:60.57ms
step:1494/1575 train_time:90515ms step_avg:60.59ms
step:1495/1575 train_time:90605ms step_avg:60.61ms
step:1496/1575 train_time:90692ms step_avg:60.62ms
step:1497/1575 train_time:90782ms step_avg:60.64ms
step:1498/1575 train_time:90866ms step_avg:60.66ms
step:1499/1575 train_time:90957ms step_avg:60.68ms
step:1500/1575 train_time:91041ms step_avg:60.69ms
step:1500/1575 val_loss:3.3013 train_time:91115ms step_avg:60.74ms
step:1501/1575 train_time:91135ms step_avg:60.72ms
step:1502/1575 train_time:91223ms step_avg:60.73ms
step:1503/1575 train_time:91318ms step_avg:60.76ms
step:1504/1575 train_time:91403ms step_avg:60.77ms
step:1505/1575 train_time:91492ms step_avg:60.79ms
step:1506/1575 train_time:91577ms step_avg:60.81ms
step:1507/1575 train_time:91665ms step_avg:60.83ms
step:1508/1575 train_time:91750ms step_avg:60.84ms
step:1509/1575 train_time:91840ms step_avg:60.86ms
step:1510/1575 train_time:91924ms step_avg:60.88ms
step:1511/1575 train_time:92012ms step_avg:60.90ms
step:1512/1575 train_time:92099ms step_avg:60.91ms
step:1513/1575 train_time:92189ms step_avg:60.93ms
step:1514/1575 train_time:92281ms step_avg:60.95ms
step:1515/1575 train_time:92369ms step_avg:60.97ms
step:1516/1575 train_time:92456ms step_avg:60.99ms
step:1517/1575 train_time:92545ms step_avg:61.01ms
step:1518/1575 train_time:92630ms step_avg:61.02ms
step:1519/1575 train_time:92719ms step_avg:61.04ms
step:1520/1575 train_time:92804ms step_avg:61.05ms
step:1521/1575 train_time:92892ms step_avg:61.07ms
step:1522/1575 train_time:92978ms step_avg:61.09ms
step:1523/1575 train_time:93067ms step_avg:61.11ms
step:1524/1575 train_time:93155ms step_avg:61.13ms
step:1525/1575 train_time:93248ms step_avg:61.15ms
step:1526/1575 train_time:93335ms step_avg:61.16ms
step:1527/1575 train_time:93427ms step_avg:61.18ms
step:1528/1575 train_time:93512ms step_avg:61.20ms
step:1529/1575 train_time:93601ms step_avg:61.22ms
step:1530/1575 train_time:93686ms step_avg:61.23ms
step:1531/1575 train_time:93775ms step_avg:61.25ms
step:1532/1575 train_time:93863ms step_avg:61.27ms
step:1533/1575 train_time:93949ms step_avg:61.28ms
step:1534/1575 train_time:94034ms step_avg:61.30ms
step:1535/1575 train_time:94124ms step_avg:61.32ms
step:1536/1575 train_time:94218ms step_avg:61.34ms
step:1537/1575 train_time:94307ms step_avg:61.36ms
step:1538/1575 train_time:94394ms step_avg:61.37ms
step:1539/1575 train_time:94485ms step_avg:61.39ms
step:1540/1575 train_time:94570ms step_avg:61.41ms
step:1541/1575 train_time:94660ms step_avg:61.43ms
step:1542/1575 train_time:94746ms step_avg:61.44ms
step:1543/1575 train_time:94835ms step_avg:61.46ms
step:1544/1575 train_time:94922ms step_avg:61.48ms
step:1545/1575 train_time:95009ms step_avg:61.49ms
step:1546/1575 train_time:95096ms step_avg:61.51ms
step:1547/1575 train_time:95187ms step_avg:61.53ms
step:1548/1575 train_time:95274ms step_avg:61.55ms
step:1549/1575 train_time:95365ms step_avg:61.57ms
step:1550/1575 train_time:95451ms step_avg:61.58ms
step:1551/1575 train_time:95542ms step_avg:61.60ms
step:1552/1575 train_time:95627ms step_avg:61.62ms
step:1553/1575 train_time:95717ms step_avg:61.63ms
step:1554/1575 train_time:95802ms step_avg:61.65ms
step:1555/1575 train_time:95891ms step_avg:61.67ms
step:1556/1575 train_time:95977ms step_avg:61.68ms
step:1557/1575 train_time:96067ms step_avg:61.70ms
step:1558/1575 train_time:96153ms step_avg:61.72ms
step:1559/1575 train_time:96245ms step_avg:61.73ms
step:1560/1575 train_time:96331ms step_avg:61.75ms
step:1561/1575 train_time:96422ms step_avg:61.77ms
step:1562/1575 train_time:96507ms step_avg:61.78ms
step:1563/1575 train_time:96598ms step_avg:61.80ms
step:1564/1575 train_time:96685ms step_avg:61.82ms
step:1565/1575 train_time:96773ms step_avg:61.84ms
step:1566/1575 train_time:96859ms step_avg:61.85ms
step:1567/1575 train_time:96949ms step_avg:61.87ms
step:1568/1575 train_time:97035ms step_avg:61.88ms
step:1569/1575 train_time:97130ms step_avg:61.91ms
step:1570/1575 train_time:97214ms step_avg:61.92ms
step:1571/1575 train_time:97303ms step_avg:61.94ms
step:1572/1575 train_time:97389ms step_avg:61.95ms
step:1573/1575 train_time:97480ms step_avg:61.97ms
step:1574/1575 train_time:97565ms step_avg:61.99ms
step:1575/1575 train_time:97657ms step_avg:62.00ms
step:1575/1575 val_loss:3.2792 train_time:97725ms step_avg:62.05ms
peak memory allocated: 31016 MiB reserved: 47000 MiB