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(3)])
        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
        
        # 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 layers
        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]
        # 012 ... 012 structure on token value embeddings by @YouJiacheng, improved on @leloykun's U-net structure
        # dropping first layer updates this to .12 ... 012
        ve = [ve[1], ve[2]] + [None] * (self.num_layers - 5) + [ve[0], ve[1], ve[2]]
        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},
            "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", "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 = 1560  # 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 03:06:40 2026       
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 570.148.08             Driver Version: 570.148.08     CUDA Version: 12.8     |
|-----------------------------------------+------------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  NVIDIA H100 80GB HBM3          On  |   00000000:61:00.0 Off |                    0 |
| N/A   35C    P0            120W /  700W |    1519MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   1  NVIDIA H100 80GB HBM3          On  |   00000000:62:00.0 Off |                    0 |
| N/A   39C    P0            121W /  700W |    1519MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   2  NVIDIA H100 80GB HBM3          On  |   00000000:63:00.0 Off |                    0 |
| N/A   41C    P0            122W /  700W |    1519MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   3  NVIDIA H100 80GB HBM3          On  |   00000000:64:00.0 Off |                    0 |
| N/A   35C    P0            120W /  700W |    1519MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   4  NVIDIA H100 80GB HBM3          On  |   00000000:6A:00.0 Off |                    0 |
| N/A   35C    P0            120W /  700W |    1519MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   5  NVIDIA H100 80GB HBM3          On  |   00000000:6B:00.0 Off |                    0 |
| N/A   42C    P0            131W /  700W |    1519MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   6  NVIDIA H100 80GB HBM3          On  |   00000000:6C:00.0 Off |                    0 |
| N/A   40C    P0            124W /  700W |    1519MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   7  NVIDIA H100 80GB HBM3          On  |   00000000:6D:00.0 Off |                    0 |
| N/A   36C    P0            119W /  700W |    1519MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
                                                                                         
+-----------------------------------------------------------------------------------------+
| Processes:                                                                              |
|  GPU   GI   CI              PID   Type   Process name                        GPU Memory |
|        ID   ID                                                               Usage      |
|=========================================================================================|
|    0   N/A  N/A          285231      C   /usr/bin/python3                       1510MiB |
|    1   N/A  N/A          285232      C   /usr/bin/python3                       1510MiB |
|    2   N/A  N/A          285233      C   /usr/bin/python3                       1510MiB |
|    3   N/A  N/A          285234      C   /usr/bin/python3                       1510MiB |
|    4   N/A  N/A          285235      C   /usr/bin/python3                       1510MiB |
|    5   N/A  N/A          285236      C   /usr/bin/python3                       1510MiB |
|    6   N/A  N/A          285237      C   /usr/bin/python3                       1510MiB |
|    7   N/A  N/A          285238      C   /usr/bin/python3                       1510MiB |
+-----------------------------------------------------------------------------------------+

====================================================================================================
Compiling model and warming up kernels (~7 minutes on first execution)
Sampling steps [0, 1, 2, 519, 520, 521, 1039, 1040, 1041, 1559, 1560, 1561] for warmup
Resetting Model
step:0/1600 val_loss:10.8293 train_time:0ms step_avg:0.03ms
step:1/1600 train_time:78ms step_avg:77.80ms
step:2/1600 train_time:100ms step_avg:49.91ms
step:3/1600 train_time:119ms step_avg:39.79ms
step:4/1600 train_time:144ms step_avg:36.07ms
step:5/1600 train_time:175ms step_avg:34.94ms
step:6/1600 train_time:291ms step_avg:48.58ms
step:7/1600 train_time:310ms step_avg:44.22ms
step:8/1600 train_time:329ms step_avg:41.14ms
step:9/1600 train_time:358ms step_avg:39.81ms
step:10/1600 train_time:395ms step_avg:39.53ms
step:11/1600 train_time:426ms step_avg:38.76ms
step:12/1600 train_time:464ms step_avg:38.65ms
step:13/1600 train_time:495ms step_avg:38.05ms
step:14/1600 train_time:532ms step_avg:37.99ms
step:15/1600 train_time:563ms step_avg:37.52ms
step:16/1600 train_time:601ms step_avg:37.53ms
step:17/1600 train_time:631ms step_avg:37.13ms
step:18/1600 train_time:669ms step_avg:37.16ms
step:19/1600 train_time:699ms step_avg:36.80ms
step:20/1600 train_time:736ms step_avg:36.81ms
step:21/1600 train_time:767ms step_avg:36.52ms
step:22/1600 train_time:805ms step_avg:36.57ms
step:23/1600 train_time:836ms step_avg:36.33ms
step:24/1600 train_time:873ms step_avg:36.37ms
step:25/1600 train_time:903ms step_avg:36.13ms
step:26/1600 train_time:940ms step_avg:36.16ms
step:27/1600 train_time:971ms step_avg:35.96ms
step:28/1600 train_time:1008ms step_avg:36.00ms
step:29/1600 train_time:1039ms step_avg:35.81ms
step:30/1600 train_time:1076ms step_avg:35.86ms
step:31/1600 train_time:1106ms step_avg:35.69ms
step:32/1600 train_time:1144ms step_avg:35.74ms
step:33/1600 train_time:1174ms step_avg:35.58ms
step:34/1600 train_time:1211ms step_avg:35.62ms
step:35/1600 train_time:1242ms step_avg:35.49ms
step:36/1600 train_time:1279ms step_avg:35.54ms
step:37/1600 train_time:1310ms step_avg:35.41ms
step:38/1600 train_time:1348ms step_avg:35.47ms
step:39/1600 train_time:1378ms step_avg:35.34ms
step:40/1600 train_time:1415ms step_avg:35.39ms
step:41/1600 train_time:1447ms step_avg:35.28ms
step:42/1600 train_time:1484ms step_avg:35.33ms
step:43/1600 train_time:1514ms step_avg:35.22ms
step:44/1600 train_time:1552ms step_avg:35.27ms
step:45/1600 train_time:1582ms step_avg:35.16ms
step:46/1600 train_time:1619ms step_avg:35.20ms
step:47/1600 train_time:1650ms step_avg:35.11ms
step:48/1600 train_time:1688ms step_avg:35.16ms
step:49/1600 train_time:1718ms step_avg:35.06ms
step:50/1600 train_time:1756ms step_avg:35.11ms
step:51/1600 train_time:1786ms step_avg:35.02ms
step:52/1600 train_time:1824ms step_avg:35.07ms
step:53/1600 train_time:1854ms step_avg:34.98ms
step:54/1600 train_time:1891ms step_avg:35.03ms
step:55/1600 train_time:1922ms step_avg:34.94ms
step:56/1600 train_time:1959ms step_avg:34.98ms
step:57/1600 train_time:1990ms step_avg:34.91ms
step:58/1600 train_time:2027ms step_avg:34.95ms
step:59/1600 train_time:2058ms step_avg:34.88ms
step:60/1600 train_time:2095ms step_avg:34.92ms
step:61/1600 train_time:2126ms step_avg:34.85ms
step:62/1600 train_time:2164ms step_avg:34.91ms
step:63/1600 train_time:2195ms step_avg:34.84ms
step:64/1600 train_time:2232ms step_avg:34.88ms
step:65/1600 train_time:2263ms step_avg:34.82ms
step:66/1600 train_time:2300ms step_avg:34.85ms
step:67/1600 train_time:2331ms step_avg:34.79ms
step:68/1600 train_time:2369ms step_avg:34.84ms
step:69/1600 train_time:2399ms step_avg:34.77ms
step:70/1600 train_time:2436ms step_avg:34.81ms
step:71/1600 train_time:2467ms step_avg:34.75ms
step:72/1600 train_time:2505ms step_avg:34.79ms
step:73/1600 train_time:2536ms step_avg:34.73ms
step:74/1600 train_time:2573ms step_avg:34.77ms
step:75/1600 train_time:2604ms step_avg:34.72ms
step:76/1600 train_time:2641ms step_avg:34.75ms
step:77/1600 train_time:2672ms step_avg:34.70ms
step:78/1600 train_time:2709ms step_avg:34.74ms
step:79/1600 train_time:2740ms step_avg:34.68ms
step:80/1600 train_time:2777ms step_avg:34.71ms
step:81/1600 train_time:2807ms step_avg:34.66ms
step:82/1600 train_time:2845ms step_avg:34.69ms
step:83/1600 train_time:2875ms step_avg:34.64ms
step:84/1600 train_time:2912ms step_avg:34.67ms
step:85/1600 train_time:2943ms step_avg:34.62ms
step:86/1600 train_time:2980ms step_avg:34.66ms
step:87/1600 train_time:3011ms step_avg:34.61ms
step:88/1600 train_time:3049ms step_avg:34.64ms
step:89/1600 train_time:3079ms step_avg:34.59ms
step:90/1600 train_time:3116ms step_avg:34.62ms
step:91/1600 train_time:3147ms step_avg:34.58ms
step:92/1600 train_time:3185ms step_avg:34.61ms
step:93/1600 train_time:3215ms step_avg:34.57ms
step:94/1600 train_time:3253ms step_avg:34.60ms
step:95/1600 train_time:3283ms step_avg:34.56ms
step:96/1600 train_time:3321ms step_avg:34.59ms
step:97/1600 train_time:3351ms step_avg:34.55ms
step:98/1600 train_time:3389ms step_avg:34.58ms
step:99/1600 train_time:3420ms step_avg:34.54ms
step:100/1600 train_time:3458ms step_avg:34.58ms
step:101/1600 train_time:3488ms step_avg:34.54ms
step:102/1600 train_time:3525ms step_avg:34.56ms
step:103/1600 train_time:3556ms step_avg:34.52ms
step:104/1600 train_time:3593ms step_avg:34.55ms
step:105/1600 train_time:3624ms step_avg:34.51ms
step:106/1600 train_time:3661ms step_avg:34.54ms
step:107/1600 train_time:3692ms step_avg:34.50ms
step:108/1600 train_time:3729ms step_avg:34.53ms
step:109/1600 train_time:3759ms step_avg:34.49ms
step:110/1600 train_time:3797ms step_avg:34.52ms
step:111/1600 train_time:3828ms step_avg:34.48ms
step:112/1600 train_time:3865ms step_avg:34.51ms
step:113/1600 train_time:3896ms step_avg:34.48ms
step:114/1600 train_time:3933ms step_avg:34.50ms
step:115/1600 train_time:3964ms step_avg:34.47ms
step:116/1600 train_time:4001ms step_avg:34.49ms
step:117/1600 train_time:4033ms step_avg:34.47ms
step:118/1600 train_time:4070ms step_avg:34.49ms
step:119/1600 train_time:4101ms step_avg:34.46ms
step:120/1600 train_time:4137ms step_avg:34.48ms
step:121/1600 train_time:4168ms step_avg:34.45ms
step:122/1600 train_time:4205ms step_avg:34.47ms
step:123/1600 train_time:4236ms step_avg:34.44ms
step:124/1600 train_time:4273ms step_avg:34.46ms
step:125/1600 train_time:4304ms step_avg:34.43ms
step:126/1600 train_time:4341ms step_avg:34.45ms
step:127/1600 train_time:4372ms step_avg:34.42ms
step:128/1600 train_time:4409ms step_avg:34.44ms
step:129/1600 train_time:4439ms step_avg:34.41ms
step:130/1600 train_time:4477ms step_avg:34.44ms
step:131/1600 train_time:4507ms step_avg:34.40ms
step:132/1600 train_time:4545ms step_avg:34.43ms
step:133/1600 train_time:4575ms step_avg:34.40ms
step:134/1600 train_time:4612ms step_avg:34.42ms
step:135/1600 train_time:4643ms step_avg:34.39ms
step:136/1600 train_time:4680ms step_avg:34.41ms
step:137/1600 train_time:4711ms step_avg:34.39ms
step:138/1600 train_time:4748ms step_avg:34.41ms
step:139/1600 train_time:4779ms step_avg:34.38ms
step:140/1600 train_time:4816ms step_avg:34.40ms
step:141/1600 train_time:4846ms step_avg:34.37ms
step:142/1600 train_time:4884ms step_avg:34.39ms
step:143/1600 train_time:4914ms step_avg:34.37ms
step:144/1600 train_time:4951ms step_avg:34.39ms
step:145/1600 train_time:4982ms step_avg:34.36ms
step:146/1600 train_time:5020ms step_avg:34.38ms
step:147/1600 train_time:5050ms step_avg:34.35ms
step:148/1600 train_time:5087ms step_avg:34.37ms
step:149/1600 train_time:5118ms step_avg:34.35ms
step:150/1600 train_time:5155ms step_avg:34.37ms
step:151/1600 train_time:5186ms step_avg:34.34ms
step:152/1600 train_time:5224ms step_avg:34.37ms
step:153/1600 train_time:5255ms step_avg:34.34ms
step:154/1600 train_time:5292ms step_avg:34.36ms
step:155/1600 train_time:5323ms step_avg:34.34ms
step:156/1600 train_time:5361ms step_avg:34.36ms
step:157/1600 train_time:5391ms step_avg:34.34ms
step:158/1600 train_time:5428ms step_avg:34.36ms
step:159/1600 train_time:5459ms step_avg:34.33ms
step:160/1600 train_time:5497ms step_avg:34.35ms
step:161/1600 train_time:5527ms step_avg:34.33ms
step:162/1600 train_time:5565ms step_avg:34.35ms
step:163/1600 train_time:5595ms step_avg:34.33ms
step:164/1600 train_time:5632ms step_avg:34.34ms
step:165/1600 train_time:5663ms step_avg:34.32ms
step:166/1600 train_time:5701ms step_avg:34.34ms
step:167/1600 train_time:5732ms step_avg:34.32ms
step:168/1600 train_time:5769ms step_avg:34.34ms
step:169/1600 train_time:5800ms step_avg:34.32ms
step:170/1600 train_time:5837ms step_avg:34.34ms
step:171/1600 train_time:5868ms step_avg:34.31ms
step:172/1600 train_time:5905ms step_avg:34.33ms
step:173/1600 train_time:5935ms step_avg:34.31ms
step:174/1600 train_time:5972ms step_avg:34.32ms
step:175/1600 train_time:6003ms step_avg:34.30ms
step:176/1600 train_time:6040ms step_avg:34.32ms
step:177/1600 train_time:6071ms step_avg:34.30ms
step:178/1600 train_time:6108ms step_avg:34.32ms
step:179/1600 train_time:6138ms step_avg:34.29ms
step:180/1600 train_time:6176ms step_avg:34.31ms
step:181/1600 train_time:6206ms step_avg:34.29ms
step:182/1600 train_time:6244ms step_avg:34.31ms
step:183/1600 train_time:6274ms step_avg:34.29ms
step:184/1600 train_time:6312ms step_avg:34.30ms
step:185/1600 train_time:6342ms step_avg:34.28ms
step:186/1600 train_time:6379ms step_avg:34.30ms
step:187/1600 train_time:6410ms step_avg:34.28ms
step:188/1600 train_time:6447ms step_avg:34.29ms
step:189/1600 train_time:6478ms step_avg:34.27ms
step:190/1600 train_time:6515ms step_avg:34.29ms
step:191/1600 train_time:6546ms step_avg:34.27ms
step:192/1600 train_time:6583ms step_avg:34.29ms
step:193/1600 train_time:6614ms step_avg:34.27ms
step:194/1600 train_time:6651ms step_avg:34.28ms
step:195/1600 train_time:6681ms step_avg:34.26ms
step:196/1600 train_time:6719ms step_avg:34.28ms
step:197/1600 train_time:6749ms step_avg:34.26ms
step:198/1600 train_time:6787ms step_avg:34.28ms
step:199/1600 train_time:6817ms step_avg:34.26ms
step:200/1600 train_time:6854ms step_avg:34.27ms
step:201/1600 train_time:6885ms step_avg:34.25ms
step:202/1600 train_time:6923ms step_avg:34.27ms
step:203/1600 train_time:6953ms step_avg:34.25ms
step:204/1600 train_time:6991ms step_avg:34.27ms
step:205/1600 train_time:7021ms step_avg:34.25ms
step:206/1600 train_time:7058ms step_avg:34.26ms
step:207/1600 train_time:7089ms step_avg:34.25ms
step:208/1600 train_time:7127ms step_avg:34.26ms
step:209/1600 train_time:7157ms step_avg:34.24ms
step:210/1600 train_time:7194ms step_avg:34.26ms
step:211/1600 train_time:7225ms step_avg:34.24ms
step:212/1600 train_time:7262ms step_avg:34.26ms
step:213/1600 train_time:7293ms step_avg:34.24ms
step:214/1600 train_time:7330ms step_avg:34.25ms
step:215/1600 train_time:7361ms step_avg:34.24ms
step:216/1600 train_time:7399ms step_avg:34.26ms
step:217/1600 train_time:7430ms step_avg:34.24ms
step:218/1600 train_time:7467ms step_avg:34.25ms
step:219/1600 train_time:7498ms step_avg:34.24ms
step:220/1600 train_time:7535ms step_avg:34.25ms
step:221/1600 train_time:7565ms step_avg:34.23ms
step:222/1600 train_time:7603ms step_avg:34.25ms
step:223/1600 train_time:7633ms step_avg:34.23ms
step:224/1600 train_time:7671ms step_avg:34.24ms
step:225/1600 train_time:7701ms step_avg:34.23ms
step:226/1600 train_time:7738ms step_avg:34.24ms
step:227/1600 train_time:7769ms step_avg:34.22ms
step:228/1600 train_time:7806ms step_avg:34.24ms
step:229/1600 train_time:7836ms step_avg:34.22ms
step:230/1600 train_time:7873ms step_avg:34.23ms
step:231/1600 train_time:7904ms step_avg:34.22ms
step:232/1600 train_time:7942ms step_avg:34.23ms
step:233/1600 train_time:7973ms step_avg:34.22ms
step:234/1600 train_time:8010ms step_avg:34.23ms
step:235/1600 train_time:8040ms step_avg:34.21ms
step:236/1600 train_time:8077ms step_avg:34.23ms
step:237/1600 train_time:8108ms step_avg:34.21ms
step:238/1600 train_time:8145ms step_avg:34.22ms
step:239/1600 train_time:8175ms step_avg:34.21ms
step:240/1600 train_time:8213ms step_avg:34.22ms
step:241/1600 train_time:8243ms step_avg:34.20ms
step:242/1600 train_time:8281ms step_avg:34.22ms
step:243/1600 train_time:8311ms step_avg:34.20ms
step:244/1600 train_time:8348ms step_avg:34.21ms
step:245/1600 train_time:8379ms step_avg:34.20ms
step:246/1600 train_time:8416ms step_avg:34.21ms
step:247/1600 train_time:8447ms step_avg:34.20ms
step:248/1600 train_time:8485ms step_avg:34.21ms
step:249/1600 train_time:8516ms step_avg:34.20ms
step:250/1600 train_time:8553ms step_avg:34.21ms
step:250/1600 val_loss:4.5848 train_time:8601ms step_avg:34.40ms
step:251/1600 train_time:8621ms step_avg:34.35ms
step:252/1600 train_time:8641ms step_avg:34.29ms
step:253/1600 train_time:8659ms step_avg:34.22ms
step:254/1600 train_time:8695ms step_avg:34.23ms
step:255/1600 train_time:8728ms step_avg:34.23ms
step:256/1600 train_time:8768ms step_avg:34.25ms
step:257/1600 train_time:8799ms step_avg:34.24ms
step:258/1600 train_time:8837ms step_avg:34.25ms
step:259/1600 train_time:8868ms step_avg:34.24ms
step:260/1600 train_time:8906ms step_avg:34.25ms
step:261/1600 train_time:8936ms step_avg:34.24ms
step:262/1600 train_time:8973ms step_avg:34.25ms
step:263/1600 train_time:9004ms step_avg:34.24ms
step:264/1600 train_time:9041ms step_avg:34.25ms
step:265/1600 train_time:9071ms step_avg:34.23ms
step:266/1600 train_time:9109ms step_avg:34.24ms
step:267/1600 train_time:9139ms step_avg:34.23ms
step:268/1600 train_time:9176ms step_avg:34.24ms
step:269/1600 train_time:9207ms step_avg:34.23ms
step:270/1600 train_time:9244ms step_avg:34.24ms
step:271/1600 train_time:9275ms step_avg:34.22ms
step:272/1600 train_time:9312ms step_avg:34.23ms
step:273/1600 train_time:9342ms step_avg:34.22ms
step:274/1600 train_time:9379ms step_avg:34.23ms
step:275/1600 train_time:9410ms step_avg:34.22ms
step:276/1600 train_time:9447ms step_avg:34.23ms
step:277/1600 train_time:9477ms step_avg:34.21ms
step:278/1600 train_time:9514ms step_avg:34.22ms
step:279/1600 train_time:9545ms step_avg:34.21ms
step:280/1600 train_time:9582ms step_avg:34.22ms
step:281/1600 train_time:9612ms step_avg:34.21ms
step:282/1600 train_time:9650ms step_avg:34.22ms
step:283/1600 train_time:9680ms step_avg:34.21ms
step:284/1600 train_time:9717ms step_avg:34.22ms
step:285/1600 train_time:9748ms step_avg:34.20ms
step:286/1600 train_time:9786ms step_avg:34.22ms
step:287/1600 train_time:9816ms step_avg:34.20ms
step:288/1600 train_time:9854ms step_avg:34.22ms
step:289/1600 train_time:9885ms step_avg:34.20ms
step:290/1600 train_time:9922ms step_avg:34.21ms
step:291/1600 train_time:9953ms step_avg:34.20ms
step:292/1600 train_time:9990ms step_avg:34.21ms
step:293/1600 train_time:10021ms step_avg:34.20ms
step:294/1600 train_time:10058ms step_avg:34.21ms
step:295/1600 train_time:10088ms step_avg:34.20ms
step:296/1600 train_time:10126ms step_avg:34.21ms
step:297/1600 train_time:10156ms step_avg:34.20ms
step:298/1600 train_time:10193ms step_avg:34.21ms
step:299/1600 train_time:10224ms step_avg:34.19ms
step:300/1600 train_time:10261ms step_avg:34.20ms
step:301/1600 train_time:10292ms step_avg:34.19ms
step:302/1600 train_time:10329ms step_avg:34.20ms
step:303/1600 train_time:10360ms step_avg:34.19ms
step:304/1600 train_time:10396ms step_avg:34.20ms
step:305/1600 train_time:10427ms step_avg:34.19ms
step:306/1600 train_time:10465ms step_avg:34.20ms
step:307/1600 train_time:10495ms step_avg:34.19ms
step:308/1600 train_time:10533ms step_avg:34.20ms
step:309/1600 train_time:10563ms step_avg:34.18ms
step:310/1600 train_time:10600ms step_avg:34.19ms
step:311/1600 train_time:10630ms step_avg:34.18ms
step:312/1600 train_time:10667ms step_avg:34.19ms
step:313/1600 train_time:10698ms step_avg:34.18ms
step:314/1600 train_time:10735ms step_avg:34.19ms
step:315/1600 train_time:10765ms step_avg:34.18ms
step:316/1600 train_time:10803ms step_avg:34.19ms
step:317/1600 train_time:10833ms step_avg:34.17ms
step:318/1600 train_time:10871ms step_avg:34.18ms
step:319/1600 train_time:10901ms step_avg:34.17ms
step:320/1600 train_time:10939ms step_avg:34.18ms
step:321/1600 train_time:10969ms step_avg:34.17ms
step:322/1600 train_time:11007ms step_avg:34.18ms
step:323/1600 train_time:11038ms step_avg:34.17ms
step:324/1600 train_time:11075ms step_avg:34.18ms
step:325/1600 train_time:11106ms step_avg:34.17ms
step:326/1600 train_time:11143ms step_avg:34.18ms
step:327/1600 train_time:11174ms step_avg:34.17ms
step:328/1600 train_time:11211ms step_avg:34.18ms
step:329/1600 train_time:11241ms step_avg:34.17ms
step:330/1600 train_time:11278ms step_avg:34.18ms
step:331/1600 train_time:11309ms step_avg:34.17ms
step:332/1600 train_time:11346ms step_avg:34.18ms
step:333/1600 train_time:11377ms step_avg:34.16ms
step:334/1600 train_time:11414ms step_avg:34.17ms
step:335/1600 train_time:11445ms step_avg:34.16ms
step:336/1600 train_time:11482ms step_avg:34.17ms
step:337/1600 train_time:11513ms step_avg:34.16ms
step:338/1600 train_time:11550ms step_avg:34.17ms
step:339/1600 train_time:11581ms step_avg:34.16ms
step:340/1600 train_time:11619ms step_avg:34.17ms
step:341/1600 train_time:11649ms step_avg:34.16ms
step:342/1600 train_time:11687ms step_avg:34.17ms
step:343/1600 train_time:11717ms step_avg:34.16ms
step:344/1600 train_time:11754ms step_avg:34.17ms
step:345/1600 train_time:11784ms step_avg:34.16ms
step:346/1600 train_time:11821ms step_avg:34.17ms
step:347/1600 train_time:11852ms step_avg:34.15ms
step:348/1600 train_time:11889ms step_avg:34.16ms
step:349/1600 train_time:11919ms step_avg:34.15ms
step:350/1600 train_time:11956ms step_avg:34.16ms
step:351/1600 train_time:11987ms step_avg:34.15ms
step:352/1600 train_time:12024ms step_avg:34.16ms
step:353/1600 train_time:12055ms step_avg:34.15ms
step:354/1600 train_time:12093ms step_avg:34.16ms
step:355/1600 train_time:12123ms step_avg:34.15ms
step:356/1600 train_time:12160ms step_avg:34.16ms
step:357/1600 train_time:12190ms step_avg:34.15ms
step:358/1600 train_time:12228ms step_avg:34.16ms
step:359/1600 train_time:12258ms step_avg:34.14ms
step:360/1600 train_time:12295ms step_avg:34.15ms
step:361/1600 train_time:12326ms step_avg:34.14ms
step:362/1600 train_time:12363ms step_avg:34.15ms
step:363/1600 train_time:12394ms step_avg:34.14ms
step:364/1600 train_time:12431ms step_avg:34.15ms
step:365/1600 train_time:12461ms step_avg:34.14ms
step:366/1600 train_time:12498ms step_avg:34.15ms
step:367/1600 train_time:12529ms step_avg:34.14ms
step:368/1600 train_time:12566ms step_avg:34.15ms
step:369/1600 train_time:12596ms step_avg:34.14ms
step:370/1600 train_time:12633ms step_avg:34.14ms
step:371/1600 train_time:12664ms step_avg:34.13ms
step:372/1600 train_time:12700ms step_avg:34.14ms
step:373/1600 train_time:12731ms step_avg:34.13ms
step:374/1600 train_time:12768ms step_avg:34.14ms
step:375/1600 train_time:12799ms step_avg:34.13ms
step:376/1600 train_time:12836ms step_avg:34.14ms
step:377/1600 train_time:12867ms step_avg:34.13ms
step:378/1600 train_time:12904ms step_avg:34.14ms
step:379/1600 train_time:12934ms step_avg:34.13ms
step:380/1600 train_time:12972ms step_avg:34.14ms
step:381/1600 train_time:13002ms step_avg:34.13ms
step:382/1600 train_time:13039ms step_avg:34.13ms
step:383/1600 train_time:13070ms step_avg:34.12ms
step:384/1600 train_time:13107ms step_avg:34.13ms
step:385/1600 train_time:13138ms step_avg:34.13ms
step:386/1600 train_time:13175ms step_avg:34.13ms
step:387/1600 train_time:13206ms step_avg:34.12ms
step:388/1600 train_time:13243ms step_avg:34.13ms
step:389/1600 train_time:13274ms step_avg:34.12ms
step:390/1600 train_time:13311ms step_avg:34.13ms
step:391/1600 train_time:13342ms step_avg:34.12ms
step:392/1600 train_time:13379ms step_avg:34.13ms
step:393/1600 train_time:13409ms step_avg:34.12ms
step:394/1600 train_time:13447ms step_avg:34.13ms
step:395/1600 train_time:13477ms step_avg:34.12ms
step:396/1600 train_time:13514ms step_avg:34.13ms
step:397/1600 train_time:13545ms step_avg:34.12ms
step:398/1600 train_time:13582ms step_avg:34.13ms
step:399/1600 train_time:13613ms step_avg:34.12ms
step:400/1600 train_time:13650ms step_avg:34.12ms
step:401/1600 train_time:13680ms step_avg:34.11ms
step:402/1600 train_time:13717ms step_avg:34.12ms
step:403/1600 train_time:13748ms step_avg:34.11ms
step:404/1600 train_time:13785ms step_avg:34.12ms
step:405/1600 train_time:13815ms step_avg:34.11ms
step:406/1600 train_time:13853ms step_avg:34.12ms
step:407/1600 train_time:13883ms step_avg:34.11ms
step:408/1600 train_time:13920ms step_avg:34.12ms
step:409/1600 train_time:13951ms step_avg:34.11ms
step:410/1600 train_time:13988ms step_avg:34.12ms
step:411/1600 train_time:14019ms step_avg:34.11ms
step:412/1600 train_time:14056ms step_avg:34.12ms
step:413/1600 train_time:14086ms step_avg:34.11ms
step:414/1600 train_time:14124ms step_avg:34.11ms
step:415/1600 train_time:14154ms step_avg:34.11ms
step:416/1600 train_time:14192ms step_avg:34.11ms
step:417/1600 train_time:14222ms step_avg:34.11ms
step:418/1600 train_time:14259ms step_avg:34.11ms
step:419/1600 train_time:14290ms step_avg:34.10ms
step:420/1600 train_time:14327ms step_avg:34.11ms
step:421/1600 train_time:14358ms step_avg:34.10ms
step:422/1600 train_time:14395ms step_avg:34.11ms
step:423/1600 train_time:14426ms step_avg:34.10ms
step:424/1600 train_time:14463ms step_avg:34.11ms
step:425/1600 train_time:14493ms step_avg:34.10ms
step:426/1600 train_time:14531ms step_avg:34.11ms
step:427/1600 train_time:14561ms step_avg:34.10ms
step:428/1600 train_time:14599ms step_avg:34.11ms
step:429/1600 train_time:14629ms step_avg:34.10ms
step:430/1600 train_time:14666ms step_avg:34.11ms
step:431/1600 train_time:14697ms step_avg:34.10ms
step:432/1600 train_time:14734ms step_avg:34.11ms
step:433/1600 train_time:14765ms step_avg:34.10ms
step:434/1600 train_time:14802ms step_avg:34.11ms
step:435/1600 train_time:14832ms step_avg:34.10ms
step:436/1600 train_time:14870ms step_avg:34.11ms
step:437/1600 train_time:14900ms step_avg:34.10ms
step:438/1600 train_time:14937ms step_avg:34.10ms
step:439/1600 train_time:14968ms step_avg:34.10ms
step:440/1600 train_time:15005ms step_avg:34.10ms
step:441/1600 train_time:15036ms step_avg:34.10ms
step:442/1600 train_time:15073ms step_avg:34.10ms
step:443/1600 train_time:15104ms step_avg:34.09ms
step:444/1600 train_time:15141ms step_avg:34.10ms
step:445/1600 train_time:15171ms step_avg:34.09ms
step:446/1600 train_time:15208ms step_avg:34.10ms
step:447/1600 train_time:15239ms step_avg:34.09ms
step:448/1600 train_time:15276ms step_avg:34.10ms
step:449/1600 train_time:15306ms step_avg:34.09ms
step:450/1600 train_time:15343ms step_avg:34.10ms
step:451/1600 train_time:15374ms step_avg:34.09ms
step:452/1600 train_time:15411ms step_avg:34.09ms
step:453/1600 train_time:15441ms step_avg:34.09ms
step:454/1600 train_time:15478ms step_avg:34.09ms
step:455/1600 train_time:15509ms step_avg:34.09ms
step:456/1600 train_time:15547ms step_avg:34.09ms
step:457/1600 train_time:15577ms step_avg:34.09ms
step:458/1600 train_time:15614ms step_avg:34.09ms
step:459/1600 train_time:15645ms step_avg:34.09ms
step:460/1600 train_time:15683ms step_avg:34.09ms
step:461/1600 train_time:15713ms step_avg:34.09ms
step:462/1600 train_time:15751ms step_avg:34.09ms
step:463/1600 train_time:15781ms step_avg:34.09ms
step:464/1600 train_time:15819ms step_avg:34.09ms
step:465/1600 train_time:15849ms step_avg:34.08ms
step:466/1600 train_time:15886ms step_avg:34.09ms
step:467/1600 train_time:15917ms step_avg:34.08ms
step:468/1600 train_time:15954ms step_avg:34.09ms
step:469/1600 train_time:15984ms step_avg:34.08ms
step:470/1600 train_time:16021ms step_avg:34.09ms
step:471/1600 train_time:16052ms step_avg:34.08ms
step:472/1600 train_time:16089ms step_avg:34.09ms
step:473/1600 train_time:16119ms step_avg:34.08ms
step:474/1600 train_time:16156ms step_avg:34.09ms
step:475/1600 train_time:16187ms step_avg:34.08ms
step:476/1600 train_time:16224ms step_avg:34.08ms
step:477/1600 train_time:16255ms step_avg:34.08ms
step:478/1600 train_time:16292ms step_avg:34.08ms
step:479/1600 train_time:16322ms step_avg:34.08ms
step:480/1600 train_time:16360ms step_avg:34.08ms
step:481/1600 train_time:16390ms step_avg:34.07ms
step:482/1600 train_time:16427ms step_avg:34.08ms
step:483/1600 train_time:16458ms step_avg:34.07ms
step:484/1600 train_time:16496ms step_avg:34.08ms
step:485/1600 train_time:16526ms step_avg:34.07ms
step:486/1600 train_time:16563ms step_avg:34.08ms
step:487/1600 train_time:16594ms step_avg:34.07ms
step:488/1600 train_time:16632ms step_avg:34.08ms
step:489/1600 train_time:16662ms step_avg:34.07ms
step:490/1600 train_time:16698ms step_avg:34.08ms
step:491/1600 train_time:16729ms step_avg:34.07ms
step:492/1600 train_time:16767ms step_avg:34.08ms
step:493/1600 train_time:16798ms step_avg:34.07ms
step:494/1600 train_time:16835ms step_avg:34.08ms
step:495/1600 train_time:16865ms step_avg:34.07ms
step:496/1600 train_time:16902ms step_avg:34.08ms
step:497/1600 train_time:16932ms step_avg:34.07ms
step:498/1600 train_time:16970ms step_avg:34.08ms
step:499/1600 train_time:17000ms step_avg:34.07ms
step:500/1600 train_time:17037ms step_avg:34.07ms
step:500/1600 val_loss:4.2420 train_time:17085ms step_avg:34.17ms
step:501/1600 train_time:17105ms step_avg:34.14ms
step:502/1600 train_time:17125ms step_avg:34.11ms
step:503/1600 train_time:17142ms step_avg:34.08ms
step:504/1600 train_time:17176ms step_avg:34.08ms
step:505/1600 train_time:17208ms step_avg:34.08ms
step:506/1600 train_time:17246ms step_avg:34.08ms
step:507/1600 train_time:17277ms step_avg:34.08ms
step:508/1600 train_time:17315ms step_avg:34.08ms
step:509/1600 train_time:17345ms step_avg:34.08ms
step:510/1600 train_time:17382ms step_avg:34.08ms
step:511/1600 train_time:17413ms step_avg:34.08ms
step:512/1600 train_time:17450ms step_avg:34.08ms
step:513/1600 train_time:17481ms step_avg:34.08ms
step:514/1600 train_time:17517ms step_avg:34.08ms
step:515/1600 train_time:17548ms step_avg:34.07ms
step:516/1600 train_time:17585ms step_avg:34.08ms
step:517/1600 train_time:17616ms step_avg:34.07ms
step:518/1600 train_time:17653ms step_avg:34.08ms
step:519/1600 train_time:17683ms step_avg:34.07ms
step:520/1600 train_time:17720ms step_avg:34.08ms
step:521/1600 train_time:17794ms step_avg:34.15ms
step:522/1600 train_time:17849ms step_avg:34.19ms
step:523/1600 train_time:17910ms step_avg:34.24ms
step:524/1600 train_time:17968ms step_avg:34.29ms
step:525/1600 train_time:18029ms step_avg:34.34ms
step:526/1600 train_time:18087ms step_avg:34.39ms
step:527/1600 train_time:18149ms step_avg:34.44ms
step:528/1600 train_time:18208ms step_avg:34.49ms
step:529/1600 train_time:18272ms step_avg:34.54ms
step:530/1600 train_time:18331ms step_avg:34.59ms
step:531/1600 train_time:18394ms step_avg:34.64ms
step:532/1600 train_time:18454ms step_avg:34.69ms
step:533/1600 train_time:18516ms step_avg:34.74ms
step:534/1600 train_time:18576ms step_avg:34.79ms
step:535/1600 train_time:18639ms step_avg:34.84ms
step:536/1600 train_time:18698ms step_avg:34.88ms
step:537/1600 train_time:18760ms step_avg:34.94ms
step:538/1600 train_time:18819ms step_avg:34.98ms
step:539/1600 train_time:18881ms step_avg:35.03ms
step:540/1600 train_time:18940ms step_avg:35.07ms
step:541/1600 train_time:19003ms step_avg:35.13ms
step:542/1600 train_time:19062ms step_avg:35.17ms
step:543/1600 train_time:19124ms step_avg:35.22ms
step:544/1600 train_time:19184ms step_avg:35.26ms
step:545/1600 train_time:19246ms step_avg:35.31ms
step:546/1600 train_time:19305ms step_avg:35.36ms
step:547/1600 train_time:19367ms step_avg:35.41ms
step:548/1600 train_time:19425ms step_avg:35.45ms
step:549/1600 train_time:19487ms step_avg:35.50ms
step:550/1600 train_time:19546ms step_avg:35.54ms
step:551/1600 train_time:19608ms step_avg:35.59ms
step:552/1600 train_time:19667ms step_avg:35.63ms
step:553/1600 train_time:19728ms step_avg:35.68ms
step:554/1600 train_time:19787ms step_avg:35.72ms
step:555/1600 train_time:19849ms step_avg:35.76ms
step:556/1600 train_time:19907ms step_avg:35.80ms
step:557/1600 train_time:19969ms step_avg:35.85ms
step:558/1600 train_time:20028ms step_avg:35.89ms
step:559/1600 train_time:20089ms step_avg:35.94ms
step:560/1600 train_time:20148ms step_avg:35.98ms
step:561/1600 train_time:20210ms step_avg:36.02ms
step:562/1600 train_time:20268ms step_avg:36.06ms
step:563/1600 train_time:20331ms step_avg:36.11ms
step:564/1600 train_time:20390ms step_avg:36.15ms
step:565/1600 train_time:20453ms step_avg:36.20ms
step:566/1600 train_time:20512ms step_avg:36.24ms
step:567/1600 train_time:20577ms step_avg:36.29ms
step:568/1600 train_time:20636ms step_avg:36.33ms
step:569/1600 train_time:20697ms step_avg:36.37ms
step:570/1600 train_time:20756ms step_avg:36.41ms
step:571/1600 train_time:20820ms step_avg:36.46ms
step:572/1600 train_time:20877ms step_avg:36.50ms
step:573/1600 train_time:20944ms step_avg:36.55ms
step:574/1600 train_time:21002ms step_avg:36.59ms
step:575/1600 train_time:21064ms step_avg:36.63ms
step:576/1600 train_time:21123ms step_avg:36.67ms
step:577/1600 train_time:21185ms step_avg:36.72ms
step:578/1600 train_time:21243ms step_avg:36.75ms
step:579/1600 train_time:21306ms step_avg:36.80ms
step:580/1600 train_time:21364ms step_avg:36.83ms
step:581/1600 train_time:21427ms step_avg:36.88ms
step:582/1600 train_time:21486ms step_avg:36.92ms
step:583/1600 train_time:21549ms step_avg:36.96ms
step:584/1600 train_time:21609ms step_avg:37.00ms
step:585/1600 train_time:21670ms step_avg:37.04ms
step:586/1600 train_time:21729ms step_avg:37.08ms
step:587/1600 train_time:21790ms step_avg:37.12ms
step:588/1600 train_time:21848ms step_avg:37.16ms
step:589/1600 train_time:21911ms step_avg:37.20ms
step:590/1600 train_time:21970ms step_avg:37.24ms
step:591/1600 train_time:22034ms step_avg:37.28ms
step:592/1600 train_time:22092ms step_avg:37.32ms
step:593/1600 train_time:22154ms step_avg:37.36ms
step:594/1600 train_time:22214ms step_avg:37.40ms
step:595/1600 train_time:22277ms step_avg:37.44ms
step:596/1600 train_time:22336ms step_avg:37.48ms
step:597/1600 train_time:22399ms step_avg:37.52ms
step:598/1600 train_time:22458ms step_avg:37.56ms
step:599/1600 train_time:22521ms step_avg:37.60ms
step:600/1600 train_time:22580ms step_avg:37.63ms
step:601/1600 train_time:22643ms step_avg:37.68ms
step:602/1600 train_time:22703ms step_avg:37.71ms
step:603/1600 train_time:22765ms step_avg:37.75ms
step:604/1600 train_time:22823ms step_avg:37.79ms
step:605/1600 train_time:22885ms step_avg:37.83ms
step:606/1600 train_time:22944ms step_avg:37.86ms
step:607/1600 train_time:23007ms step_avg:37.90ms
step:608/1600 train_time:23065ms step_avg:37.94ms
step:609/1600 train_time:23127ms step_avg:37.98ms
step:610/1600 train_time:23185ms step_avg:38.01ms
step:611/1600 train_time:23248ms step_avg:38.05ms
step:612/1600 train_time:23306ms step_avg:38.08ms
step:613/1600 train_time:23368ms step_avg:38.12ms
step:614/1600 train_time:23427ms step_avg:38.15ms
step:615/1600 train_time:23488ms step_avg:38.19ms
step:616/1600 train_time:23548ms step_avg:38.23ms
step:617/1600 train_time:23609ms step_avg:38.26ms
step:618/1600 train_time:23668ms step_avg:38.30ms
step:619/1600 train_time:23730ms step_avg:38.34ms
step:620/1600 train_time:23789ms step_avg:38.37ms
step:621/1600 train_time:23851ms step_avg:38.41ms
step:622/1600 train_time:23910ms step_avg:38.44ms
step:623/1600 train_time:23973ms step_avg:38.48ms
step:624/1600 train_time:24032ms step_avg:38.51ms
step:625/1600 train_time:24094ms step_avg:38.55ms
step:626/1600 train_time:24154ms step_avg:38.58ms
step:627/1600 train_time:24216ms step_avg:38.62ms
step:628/1600 train_time:24275ms step_avg:38.65ms
step:629/1600 train_time:24338ms step_avg:38.69ms
step:630/1600 train_time:24397ms step_avg:38.72ms
step:631/1600 train_time:24459ms step_avg:38.76ms
step:632/1600 train_time:24518ms step_avg:38.79ms
step:633/1600 train_time:24582ms step_avg:38.83ms
step:634/1600 train_time:24640ms step_avg:38.87ms
step:635/1600 train_time:24702ms step_avg:38.90ms
step:636/1600 train_time:24761ms step_avg:38.93ms
step:637/1600 train_time:24823ms step_avg:38.97ms
step:638/1600 train_time:24882ms step_avg:39.00ms
step:639/1600 train_time:24945ms step_avg:39.04ms
step:640/1600 train_time:25005ms step_avg:39.07ms
step:641/1600 train_time:25067ms step_avg:39.11ms
step:642/1600 train_time:25125ms step_avg:39.14ms
step:643/1600 train_time:25188ms step_avg:39.17ms
step:644/1600 train_time:25246ms step_avg:39.20ms
step:645/1600 train_time:25308ms step_avg:39.24ms
step:646/1600 train_time:25367ms step_avg:39.27ms
step:647/1600 train_time:25428ms step_avg:39.30ms
step:648/1600 train_time:25486ms step_avg:39.33ms
step:649/1600 train_time:25549ms step_avg:39.37ms
step:650/1600 train_time:25607ms step_avg:39.39ms
step:651/1600 train_time:25671ms step_avg:39.43ms
step:652/1600 train_time:25730ms step_avg:39.46ms
step:653/1600 train_time:25793ms step_avg:39.50ms
step:654/1600 train_time:25850ms step_avg:39.53ms
step:655/1600 train_time:25915ms step_avg:39.56ms
step:656/1600 train_time:25973ms step_avg:39.59ms
step:657/1600 train_time:26036ms step_avg:39.63ms
step:658/1600 train_time:26094ms step_avg:39.66ms
step:659/1600 train_time:26156ms step_avg:39.69ms
step:660/1600 train_time:26216ms step_avg:39.72ms
step:661/1600 train_time:26277ms step_avg:39.75ms
step:662/1600 train_time:26337ms step_avg:39.78ms
step:663/1600 train_time:26399ms step_avg:39.82ms
step:664/1600 train_time:26458ms step_avg:39.85ms
step:665/1600 train_time:26520ms step_avg:39.88ms
step:666/1600 train_time:26579ms step_avg:39.91ms
step:667/1600 train_time:26641ms step_avg:39.94ms
step:668/1600 train_time:26703ms step_avg:39.98ms
step:669/1600 train_time:26764ms step_avg:40.01ms
step:670/1600 train_time:26823ms step_avg:40.03ms
step:671/1600 train_time:26886ms step_avg:40.07ms
step:672/1600 train_time:26945ms step_avg:40.10ms
step:673/1600 train_time:27007ms step_avg:40.13ms
step:674/1600 train_time:27066ms step_avg:40.16ms
step:675/1600 train_time:27127ms step_avg:40.19ms
step:676/1600 train_time:27185ms step_avg:40.21ms
step:677/1600 train_time:27248ms step_avg:40.25ms
step:678/1600 train_time:27307ms step_avg:40.28ms
step:679/1600 train_time:27368ms step_avg:40.31ms
step:680/1600 train_time:27427ms step_avg:40.33ms
step:681/1600 train_time:27487ms step_avg:40.36ms
step:682/1600 train_time:27548ms step_avg:40.39ms
step:683/1600 train_time:27608ms step_avg:40.42ms
step:684/1600 train_time:27667ms step_avg:40.45ms
step:685/1600 train_time:27729ms step_avg:40.48ms
step:686/1600 train_time:27787ms step_avg:40.51ms
step:687/1600 train_time:27849ms step_avg:40.54ms
step:688/1600 train_time:27908ms step_avg:40.56ms
step:689/1600 train_time:27971ms step_avg:40.60ms
step:690/1600 train_time:28029ms step_avg:40.62ms
step:691/1600 train_time:28092ms step_avg:40.65ms
step:692/1600 train_time:28151ms step_avg:40.68ms
step:693/1600 train_time:28214ms step_avg:40.71ms
step:694/1600 train_time:28273ms step_avg:40.74ms
step:695/1600 train_time:28335ms step_avg:40.77ms
step:696/1600 train_time:28394ms step_avg:40.80ms
step:697/1600 train_time:28457ms step_avg:40.83ms
step:698/1600 train_time:28516ms step_avg:40.85ms
step:699/1600 train_time:28579ms step_avg:40.89ms
step:700/1600 train_time:28638ms step_avg:40.91ms
step:701/1600 train_time:28701ms step_avg:40.94ms
step:702/1600 train_time:28760ms step_avg:40.97ms
step:703/1600 train_time:28823ms step_avg:41.00ms
step:704/1600 train_time:28881ms step_avg:41.02ms
step:705/1600 train_time:28944ms step_avg:41.06ms
step:706/1600 train_time:29004ms step_avg:41.08ms
step:707/1600 train_time:29066ms step_avg:41.11ms
step:708/1600 train_time:29125ms step_avg:41.14ms
step:709/1600 train_time:29187ms step_avg:41.17ms
step:710/1600 train_time:29246ms step_avg:41.19ms
step:711/1600 train_time:29309ms step_avg:41.22ms
step:712/1600 train_time:29367ms step_avg:41.25ms
step:713/1600 train_time:29428ms step_avg:41.27ms
step:714/1600 train_time:29486ms step_avg:41.30ms
step:715/1600 train_time:29549ms step_avg:41.33ms
step:716/1600 train_time:29606ms step_avg:41.35ms
step:717/1600 train_time:29668ms step_avg:41.38ms
step:718/1600 train_time:29727ms step_avg:41.40ms
step:719/1600 train_time:29788ms step_avg:41.43ms
step:720/1600 train_time:29848ms step_avg:41.46ms
step:721/1600 train_time:29909ms step_avg:41.48ms
step:722/1600 train_time:29969ms step_avg:41.51ms
step:723/1600 train_time:30031ms step_avg:41.54ms
step:724/1600 train_time:30090ms step_avg:41.56ms
step:725/1600 train_time:30153ms step_avg:41.59ms
step:726/1600 train_time:30212ms step_avg:41.61ms
step:727/1600 train_time:30274ms step_avg:41.64ms
step:728/1600 train_time:30333ms step_avg:41.67ms
step:729/1600 train_time:30394ms step_avg:41.69ms
step:730/1600 train_time:30454ms step_avg:41.72ms
step:731/1600 train_time:30516ms step_avg:41.75ms
step:732/1600 train_time:30575ms step_avg:41.77ms
step:733/1600 train_time:30637ms step_avg:41.80ms
step:734/1600 train_time:30697ms step_avg:41.82ms
step:735/1600 train_time:30760ms step_avg:41.85ms
step:736/1600 train_time:30821ms step_avg:41.88ms
step:737/1600 train_time:30883ms step_avg:41.90ms
step:738/1600 train_time:30942ms step_avg:41.93ms
step:739/1600 train_time:31005ms step_avg:41.95ms
step:740/1600 train_time:31063ms step_avg:41.98ms
step:741/1600 train_time:31125ms step_avg:42.00ms
step:742/1600 train_time:31184ms step_avg:42.03ms
step:743/1600 train_time:31246ms step_avg:42.05ms
step:744/1600 train_time:31305ms step_avg:42.08ms
step:745/1600 train_time:31367ms step_avg:42.10ms
step:746/1600 train_time:31425ms step_avg:42.13ms
step:747/1600 train_time:31488ms step_avg:42.15ms
step:748/1600 train_time:31547ms step_avg:42.18ms
step:749/1600 train_time:31609ms step_avg:42.20ms
step:750/1600 train_time:31668ms step_avg:42.22ms
step:750/1600 val_loss:3.8936 train_time:31712ms step_avg:42.28ms
step:751/1600 train_time:31733ms step_avg:42.25ms
step:752/1600 train_time:31793ms step_avg:42.28ms
step:753/1600 train_time:31855ms step_avg:42.30ms
step:754/1600 train_time:31916ms step_avg:42.33ms
step:755/1600 train_time:31981ms step_avg:42.36ms
step:756/1600 train_time:32040ms step_avg:42.38ms
step:757/1600 train_time:32102ms step_avg:42.41ms
step:758/1600 train_time:32160ms step_avg:42.43ms
step:759/1600 train_time:32222ms step_avg:42.45ms
step:760/1600 train_time:32281ms step_avg:42.47ms
step:761/1600 train_time:32342ms step_avg:42.50ms
step:762/1600 train_time:32401ms step_avg:42.52ms
step:763/1600 train_time:32463ms step_avg:42.55ms
step:764/1600 train_time:32523ms step_avg:42.57ms
step:765/1600 train_time:32586ms step_avg:42.60ms
step:766/1600 train_time:32645ms step_avg:42.62ms
step:767/1600 train_time:32709ms step_avg:42.65ms
step:768/1600 train_time:32769ms step_avg:42.67ms
step:769/1600 train_time:32832ms step_avg:42.69ms
step:770/1600 train_time:32892ms step_avg:42.72ms
step:771/1600 train_time:32953ms step_avg:42.74ms
step:772/1600 train_time:33012ms step_avg:42.76ms
step:773/1600 train_time:33076ms step_avg:42.79ms
step:774/1600 train_time:33135ms step_avg:42.81ms
step:775/1600 train_time:33197ms step_avg:42.83ms
step:776/1600 train_time:33255ms step_avg:42.85ms
step:777/1600 train_time:33317ms step_avg:42.88ms
step:778/1600 train_time:33376ms step_avg:42.90ms
step:779/1600 train_time:33439ms step_avg:42.93ms
step:780/1600 train_time:33498ms step_avg:42.95ms
step:781/1600 train_time:33561ms step_avg:42.97ms
step:782/1600 train_time:33620ms step_avg:42.99ms
step:783/1600 train_time:33684ms step_avg:43.02ms
step:784/1600 train_time:33743ms step_avg:43.04ms
step:785/1600 train_time:33805ms step_avg:43.06ms
step:786/1600 train_time:33865ms step_avg:43.08ms
step:787/1600 train_time:33927ms step_avg:43.11ms
step:788/1600 train_time:33986ms step_avg:43.13ms
step:789/1600 train_time:34049ms step_avg:43.16ms
step:790/1600 train_time:34109ms step_avg:43.18ms
step:791/1600 train_time:34171ms step_avg:43.20ms
step:792/1600 train_time:34228ms step_avg:43.22ms
step:793/1600 train_time:34290ms step_avg:43.24ms
step:794/1600 train_time:34349ms step_avg:43.26ms
step:795/1600 train_time:34410ms step_avg:43.28ms
step:796/1600 train_time:34469ms step_avg:43.30ms
step:797/1600 train_time:34531ms step_avg:43.33ms
step:798/1600 train_time:34590ms step_avg:43.35ms
step:799/1600 train_time:34652ms step_avg:43.37ms
step:800/1600 train_time:34710ms step_avg:43.39ms
step:801/1600 train_time:34771ms step_avg:43.41ms
step:802/1600 train_time:34830ms step_avg:43.43ms
step:803/1600 train_time:34892ms step_avg:43.45ms
step:804/1600 train_time:34951ms step_avg:43.47ms
step:805/1600 train_time:35014ms step_avg:43.50ms
step:806/1600 train_time:35075ms step_avg:43.52ms
step:807/1600 train_time:35137ms step_avg:43.54ms
step:808/1600 train_time:35196ms step_avg:43.56ms
step:809/1600 train_time:35257ms step_avg:43.58ms
step:810/1600 train_time:35317ms step_avg:43.60ms
step:811/1600 train_time:35379ms step_avg:43.62ms
step:812/1600 train_time:35438ms step_avg:43.64ms
step:813/1600 train_time:35500ms step_avg:43.67ms
step:814/1600 train_time:35559ms step_avg:43.68ms
step:815/1600 train_time:35621ms step_avg:43.71ms
step:816/1600 train_time:35681ms step_avg:43.73ms
step:817/1600 train_time:35742ms step_avg:43.75ms
step:818/1600 train_time:35802ms step_avg:43.77ms
step:819/1600 train_time:35863ms step_avg:43.79ms
step:820/1600 train_time:35923ms step_avg:43.81ms
step:821/1600 train_time:35987ms step_avg:43.83ms
step:822/1600 train_time:36046ms step_avg:43.85ms
step:823/1600 train_time:36109ms step_avg:43.87ms
step:824/1600 train_time:36167ms step_avg:43.89ms
step:825/1600 train_time:36230ms step_avg:43.92ms
step:826/1600 train_time:36290ms step_avg:43.93ms
step:827/1600 train_time:36351ms step_avg:43.95ms
step:828/1600 train_time:36409ms step_avg:43.97ms
step:829/1600 train_time:36471ms step_avg:43.99ms
step:830/1600 train_time:36529ms step_avg:44.01ms
step:831/1600 train_time:36591ms step_avg:44.03ms
step:832/1600 train_time:36650ms step_avg:44.05ms
step:833/1600 train_time:36711ms step_avg:44.07ms
step:834/1600 train_time:36770ms step_avg:44.09ms
step:835/1600 train_time:36831ms step_avg:44.11ms
step:836/1600 train_time:36890ms step_avg:44.13ms
step:837/1600 train_time:36953ms step_avg:44.15ms
step:838/1600 train_time:37011ms step_avg:44.17ms
step:839/1600 train_time:37074ms step_avg:44.19ms
step:840/1600 train_time:37132ms step_avg:44.20ms
step:841/1600 train_time:37195ms step_avg:44.23ms
step:842/1600 train_time:37255ms step_avg:44.25ms
step:843/1600 train_time:37316ms step_avg:44.27ms
step:844/1600 train_time:37375ms step_avg:44.28ms
step:845/1600 train_time:37437ms step_avg:44.30ms
step:846/1600 train_time:37496ms step_avg:44.32ms
step:847/1600 train_time:37560ms step_avg:44.34ms
step:848/1600 train_time:37618ms step_avg:44.36ms
step:849/1600 train_time:37682ms step_avg:44.38ms
step:850/1600 train_time:37741ms step_avg:44.40ms
step:851/1600 train_time:37802ms step_avg:44.42ms
step:852/1600 train_time:37861ms step_avg:44.44ms
step:853/1600 train_time:37924ms step_avg:44.46ms
step:854/1600 train_time:37984ms step_avg:44.48ms
step:855/1600 train_time:38046ms step_avg:44.50ms
step:856/1600 train_time:38105ms step_avg:44.51ms
step:857/1600 train_time:38168ms step_avg:44.54ms
step:858/1600 train_time:38227ms step_avg:44.55ms
step:859/1600 train_time:38290ms step_avg:44.57ms
step:860/1600 train_time:38348ms step_avg:44.59ms
step:861/1600 train_time:38410ms step_avg:44.61ms
step:862/1600 train_time:38469ms step_avg:44.63ms
step:863/1600 train_time:38530ms step_avg:44.65ms
step:864/1600 train_time:38590ms step_avg:44.66ms
step:865/1600 train_time:38650ms step_avg:44.68ms
step:866/1600 train_time:38709ms step_avg:44.70ms
step:867/1600 train_time:38770ms step_avg:44.72ms
step:868/1600 train_time:38828ms step_avg:44.73ms
step:869/1600 train_time:38890ms step_avg:44.75ms
step:870/1600 train_time:38949ms step_avg:44.77ms
step:871/1600 train_time:39011ms step_avg:44.79ms
step:872/1600 train_time:39071ms step_avg:44.81ms
step:873/1600 train_time:39133ms step_avg:44.83ms
step:874/1600 train_time:39192ms step_avg:44.84ms
step:875/1600 train_time:39254ms step_avg:44.86ms
step:876/1600 train_time:39313ms step_avg:44.88ms
step:877/1600 train_time:39376ms step_avg:44.90ms
step:878/1600 train_time:39435ms step_avg:44.91ms
step:879/1600 train_time:39498ms step_avg:44.93ms
step:880/1600 train_time:39556ms step_avg:44.95ms
step:881/1600 train_time:39619ms step_avg:44.97ms
step:882/1600 train_time:39680ms step_avg:44.99ms
step:883/1600 train_time:39742ms step_avg:45.01ms
step:884/1600 train_time:39801ms step_avg:45.02ms
step:885/1600 train_time:39863ms step_avg:45.04ms
step:886/1600 train_time:39922ms step_avg:45.06ms
step:887/1600 train_time:39986ms step_avg:45.08ms
step:888/1600 train_time:40045ms step_avg:45.10ms
step:889/1600 train_time:40107ms step_avg:45.11ms
step:890/1600 train_time:40172ms step_avg:45.14ms
step:891/1600 train_time:40232ms step_avg:45.15ms
step:892/1600 train_time:40289ms step_avg:45.17ms
step:893/1600 train_time:40350ms step_avg:45.18ms
step:894/1600 train_time:40409ms step_avg:45.20ms
step:895/1600 train_time:40470ms step_avg:45.22ms
step:896/1600 train_time:40530ms step_avg:45.23ms
step:897/1600 train_time:40594ms step_avg:45.25ms
step:898/1600 train_time:40652ms step_avg:45.27ms
step:899/1600 train_time:40712ms step_avg:45.29ms
step:900/1600 train_time:40771ms step_avg:45.30ms
step:901/1600 train_time:40832ms step_avg:45.32ms
step:902/1600 train_time:40891ms step_avg:45.33ms
step:903/1600 train_time:40954ms step_avg:45.35ms
step:904/1600 train_time:41013ms step_avg:45.37ms
step:905/1600 train_time:41075ms step_avg:45.39ms
step:906/1600 train_time:41135ms step_avg:45.40ms
step:907/1600 train_time:41198ms step_avg:45.42ms
step:908/1600 train_time:41257ms step_avg:45.44ms
step:909/1600 train_time:41319ms step_avg:45.46ms
step:910/1600 train_time:41379ms step_avg:45.47ms
step:911/1600 train_time:41441ms step_avg:45.49ms
step:912/1600 train_time:41500ms step_avg:45.50ms
step:913/1600 train_time:41563ms step_avg:45.52ms
step:914/1600 train_time:41622ms step_avg:45.54ms
step:915/1600 train_time:41685ms step_avg:45.56ms
step:916/1600 train_time:41743ms step_avg:45.57ms
step:917/1600 train_time:41806ms step_avg:45.59ms
step:918/1600 train_time:41866ms step_avg:45.61ms
step:919/1600 train_time:41928ms step_avg:45.62ms
step:920/1600 train_time:41986ms step_avg:45.64ms
step:921/1600 train_time:42048ms step_avg:45.65ms
step:922/1600 train_time:42107ms step_avg:45.67ms
step:923/1600 train_time:42169ms step_avg:45.69ms
step:924/1600 train_time:42228ms step_avg:45.70ms
step:925/1600 train_time:42290ms step_avg:45.72ms
step:926/1600 train_time:42349ms step_avg:45.73ms
step:927/1600 train_time:42412ms step_avg:45.75ms
step:928/1600 train_time:42470ms step_avg:45.76ms
step:929/1600 train_time:42532ms step_avg:45.78ms
step:930/1600 train_time:42590ms step_avg:45.80ms
step:931/1600 train_time:42655ms step_avg:45.82ms
step:932/1600 train_time:42712ms step_avg:45.83ms
step:933/1600 train_time:42773ms step_avg:45.84ms
step:934/1600 train_time:42831ms step_avg:45.86ms
step:935/1600 train_time:42893ms step_avg:45.87ms
step:936/1600 train_time:42952ms step_avg:45.89ms
step:937/1600 train_time:43014ms step_avg:45.91ms
step:938/1600 train_time:43073ms step_avg:45.92ms
step:939/1600 train_time:43136ms step_avg:45.94ms
step:940/1600 train_time:43196ms step_avg:45.95ms
step:941/1600 train_time:43259ms step_avg:45.97ms
step:942/1600 train_time:43317ms step_avg:45.98ms
step:943/1600 train_time:43380ms step_avg:46.00ms
step:944/1600 train_time:43439ms step_avg:46.02ms
step:945/1600 train_time:43502ms step_avg:46.03ms
step:946/1600 train_time:43561ms step_avg:46.05ms
step:947/1600 train_time:43624ms step_avg:46.07ms
step:948/1600 train_time:43683ms step_avg:46.08ms
step:949/1600 train_time:43744ms step_avg:46.09ms
step:950/1600 train_time:43803ms step_avg:46.11ms
step:951/1600 train_time:43866ms step_avg:46.13ms
step:952/1600 train_time:43924ms step_avg:46.14ms
step:953/1600 train_time:43986ms step_avg:46.16ms
step:954/1600 train_time:44046ms step_avg:46.17ms
step:955/1600 train_time:44109ms step_avg:46.19ms
step:956/1600 train_time:44169ms step_avg:46.20ms
step:957/1600 train_time:44231ms step_avg:46.22ms
step:958/1600 train_time:44289ms step_avg:46.23ms
step:959/1600 train_time:44351ms step_avg:46.25ms
step:960/1600 train_time:44410ms step_avg:46.26ms
step:961/1600 train_time:44471ms step_avg:46.28ms
step:962/1600 train_time:44530ms step_avg:46.29ms
step:963/1600 train_time:44592ms step_avg:46.30ms
step:964/1600 train_time:44650ms step_avg:46.32ms
step:965/1600 train_time:44712ms step_avg:46.33ms
step:966/1600 train_time:44772ms step_avg:46.35ms
step:967/1600 train_time:44834ms step_avg:46.36ms
step:968/1600 train_time:44891ms step_avg:46.38ms
step:969/1600 train_time:44955ms step_avg:46.39ms
step:970/1600 train_time:45013ms step_avg:46.40ms
step:971/1600 train_time:45075ms step_avg:46.42ms
step:972/1600 train_time:45134ms step_avg:46.43ms
step:973/1600 train_time:45196ms step_avg:46.45ms
step:974/1600 train_time:45256ms step_avg:46.46ms
step:975/1600 train_time:45320ms step_avg:46.48ms
step:976/1600 train_time:45379ms step_avg:46.50ms
step:977/1600 train_time:45439ms step_avg:46.51ms
step:978/1600 train_time:45498ms step_avg:46.52ms
step:979/1600 train_time:45560ms step_avg:46.54ms
step:980/1600 train_time:45620ms step_avg:46.55ms
step:981/1600 train_time:45684ms step_avg:46.57ms
step:982/1600 train_time:45742ms step_avg:46.58ms
step:983/1600 train_time:45805ms step_avg:46.60ms
step:984/1600 train_time:45864ms step_avg:46.61ms
step:985/1600 train_time:45926ms step_avg:46.63ms
step:986/1600 train_time:45986ms step_avg:46.64ms
step:987/1600 train_time:46048ms step_avg:46.65ms
step:988/1600 train_time:46107ms step_avg:46.67ms
step:989/1600 train_time:46169ms step_avg:46.68ms
step:990/1600 train_time:46229ms step_avg:46.70ms
step:991/1600 train_time:46291ms step_avg:46.71ms
step:992/1600 train_time:46351ms step_avg:46.72ms
step:993/1600 train_time:46413ms step_avg:46.74ms
step:994/1600 train_time:46471ms step_avg:46.75ms
step:995/1600 train_time:46532ms step_avg:46.77ms
step:996/1600 train_time:46591ms step_avg:46.78ms
step:997/1600 train_time:46653ms step_avg:46.79ms
step:998/1600 train_time:46711ms step_avg:46.80ms
step:999/1600 train_time:46773ms step_avg:46.82ms
step:1000/1600 train_time:46832ms step_avg:46.83ms
step:1000/1600 val_loss:3.5958 train_time:46877ms step_avg:46.88ms
step:1001/1600 train_time:46898ms step_avg:46.85ms
step:1002/1600 train_time:46959ms step_avg:46.87ms
step:1003/1600 train_time:47021ms step_avg:46.88ms
step:1004/1600 train_time:47081ms step_avg:46.89ms
step:1005/1600 train_time:47143ms step_avg:46.91ms
step:1006/1600 train_time:47202ms step_avg:46.92ms
step:1007/1600 train_time:47264ms step_avg:46.94ms
step:1008/1600 train_time:47323ms step_avg:46.95ms
step:1009/1600 train_time:47384ms step_avg:46.96ms
step:1010/1600 train_time:47442ms step_avg:46.97ms
step:1011/1600 train_time:47506ms step_avg:46.99ms
step:1012/1600 train_time:47564ms step_avg:47.00ms
step:1013/1600 train_time:47626ms step_avg:47.01ms
step:1014/1600 train_time:47685ms step_avg:47.03ms
step:1015/1600 train_time:47747ms step_avg:47.04ms
step:1016/1600 train_time:47806ms step_avg:47.05ms
step:1017/1600 train_time:47871ms step_avg:47.07ms
step:1018/1600 train_time:47931ms step_avg:47.08ms
step:1019/1600 train_time:47994ms step_avg:47.10ms
step:1020/1600 train_time:48054ms step_avg:47.11ms
step:1021/1600 train_time:48116ms step_avg:47.13ms
step:1022/1600 train_time:48175ms step_avg:47.14ms
step:1023/1600 train_time:48238ms step_avg:47.15ms
step:1024/1600 train_time:48296ms step_avg:47.16ms
step:1025/1600 train_time:48358ms step_avg:47.18ms
step:1026/1600 train_time:48417ms step_avg:47.19ms
step:1027/1600 train_time:48479ms step_avg:47.20ms
step:1028/1600 train_time:48538ms step_avg:47.22ms
step:1029/1600 train_time:48600ms step_avg:47.23ms
step:1030/1600 train_time:48658ms step_avg:47.24ms
step:1031/1600 train_time:48719ms step_avg:47.25ms
step:1032/1600 train_time:48778ms step_avg:47.27ms
step:1033/1600 train_time:48840ms step_avg:47.28ms
step:1034/1600 train_time:48899ms step_avg:47.29ms
step:1035/1600 train_time:48961ms step_avg:47.31ms
step:1036/1600 train_time:49022ms step_avg:47.32ms
step:1037/1600 train_time:49084ms step_avg:47.33ms
step:1038/1600 train_time:49145ms step_avg:47.35ms
step:1039/1600 train_time:49206ms step_avg:47.36ms
step:1040/1600 train_time:49265ms step_avg:47.37ms
step:1041/1600 train_time:49337ms step_avg:47.39ms
step:1042/1600 train_time:49421ms step_avg:47.43ms
step:1043/1600 train_time:49509ms step_avg:47.47ms
step:1044/1600 train_time:49596ms step_avg:47.51ms
step:1045/1600 train_time:49681ms step_avg:47.54ms
step:1046/1600 train_time:49765ms step_avg:47.58ms
step:1047/1600 train_time:49853ms step_avg:47.62ms
step:1048/1600 train_time:49939ms step_avg:47.65ms
step:1049/1600 train_time:50028ms step_avg:47.69ms
step:1050/1600 train_time:50114ms step_avg:47.73ms
step:1051/1600 train_time:50201ms step_avg:47.77ms
step:1052/1600 train_time:50287ms step_avg:47.80ms
step:1053/1600 train_time:50374ms step_avg:47.84ms
step:1054/1600 train_time:50460ms step_avg:47.87ms
step:1055/1600 train_time:50548ms step_avg:47.91ms
step:1056/1600 train_time:50632ms step_avg:47.95ms
step:1057/1600 train_time:50720ms step_avg:47.99ms
step:1058/1600 train_time:50804ms step_avg:48.02ms
step:1059/1600 train_time:50892ms step_avg:48.06ms
step:1060/1600 train_time:50979ms step_avg:48.09ms
step:1061/1600 train_time:51066ms step_avg:48.13ms
step:1062/1600 train_time:51151ms step_avg:48.16ms
step:1063/1600 train_time:51240ms step_avg:48.20ms
step:1064/1600 train_time:51326ms step_avg:48.24ms
step:1065/1600 train_time:51413ms step_avg:48.27ms
step:1066/1600 train_time:51497ms step_avg:48.31ms
step:1067/1600 train_time:51587ms step_avg:48.35ms
step:1068/1600 train_time:51670ms step_avg:48.38ms
step:1069/1600 train_time:51758ms step_avg:48.42ms
step:1070/1600 train_time:51842ms step_avg:48.45ms
step:1071/1600 train_time:51931ms step_avg:48.49ms
step:1072/1600 train_time:52016ms step_avg:48.52ms
step:1073/1600 train_time:52105ms step_avg:48.56ms
step:1074/1600 train_time:52190ms step_avg:48.59ms
step:1075/1600 train_time:52278ms step_avg:48.63ms
step:1076/1600 train_time:52362ms step_avg:48.66ms
step:1077/1600 train_time:52451ms step_avg:48.70ms
step:1078/1600 train_time:52536ms step_avg:48.73ms
step:1079/1600 train_time:52624ms step_avg:48.77ms
step:1080/1600 train_time:52709ms step_avg:48.80ms
step:1081/1600 train_time:52796ms step_avg:48.84ms
step:1082/1600 train_time:52881ms step_avg:48.87ms
step:1083/1600 train_time:52969ms step_avg:48.91ms
step:1084/1600 train_time:53054ms step_avg:48.94ms
step:1085/1600 train_time:53143ms step_avg:48.98ms
step:1086/1600 train_time:53228ms step_avg:49.01ms
step:1087/1600 train_time:53317ms step_avg:49.05ms
step:1088/1600 train_time:53401ms step_avg:49.08ms
step:1089/1600 train_time:53488ms step_avg:49.12ms
step:1090/1600 train_time:53574ms step_avg:49.15ms
step:1091/1600 train_time:53663ms step_avg:49.19ms
step:1092/1600 train_time:53747ms step_avg:49.22ms
step:1093/1600 train_time:53836ms step_avg:49.26ms
step:1094/1600 train_time:53921ms step_avg:49.29ms
step:1095/1600 train_time:54009ms step_avg:49.32ms
step:1096/1600 train_time:54094ms step_avg:49.36ms
step:1097/1600 train_time:54182ms step_avg:49.39ms
step:1098/1600 train_time:54266ms step_avg:49.42ms
step:1099/1600 train_time:54354ms step_avg:49.46ms
step:1100/1600 train_time:54440ms step_avg:49.49ms
step:1101/1600 train_time:54528ms step_avg:49.53ms
step:1102/1600 train_time:54612ms step_avg:49.56ms
step:1103/1600 train_time:54702ms step_avg:49.59ms
step:1104/1600 train_time:54785ms step_avg:49.62ms
step:1105/1600 train_time:54873ms step_avg:49.66ms
step:1106/1600 train_time:54958ms step_avg:49.69ms
step:1107/1600 train_time:55047ms step_avg:49.73ms
step:1108/1600 train_time:55132ms step_avg:49.76ms
step:1109/1600 train_time:55220ms step_avg:49.79ms
step:1110/1600 train_time:55305ms step_avg:49.82ms
step:1111/1600 train_time:55392ms step_avg:49.86ms
step:1112/1600 train_time:55478ms step_avg:49.89ms
step:1113/1600 train_time:55566ms step_avg:49.92ms
step:1114/1600 train_time:55651ms step_avg:49.96ms
step:1115/1600 train_time:55740ms step_avg:49.99ms
step:1116/1600 train_time:55826ms step_avg:50.02ms
step:1117/1600 train_time:55913ms step_avg:50.06ms
step:1118/1600 train_time:56004ms step_avg:50.09ms
step:1119/1600 train_time:56089ms step_avg:50.12ms
step:1120/1600 train_time:56174ms step_avg:50.16ms
step:1121/1600 train_time:56262ms step_avg:50.19ms
step:1122/1600 train_time:56346ms step_avg:50.22ms
step:1123/1600 train_time:56434ms step_avg:50.25ms
step:1124/1600 train_time:56520ms step_avg:50.28ms
step:1125/1600 train_time:56608ms step_avg:50.32ms
step:1126/1600 train_time:56693ms step_avg:50.35ms
step:1127/1600 train_time:56781ms step_avg:50.38ms
step:1128/1600 train_time:56866ms step_avg:50.41ms
step:1129/1600 train_time:56954ms step_avg:50.45ms
step:1130/1600 train_time:57039ms step_avg:50.48ms
step:1131/1600 train_time:57127ms step_avg:50.51ms
step:1132/1600 train_time:57213ms step_avg:50.54ms
step:1133/1600 train_time:57301ms step_avg:50.57ms
step:1134/1600 train_time:57386ms step_avg:50.61ms
step:1135/1600 train_time:57474ms step_avg:50.64ms
step:1136/1600 train_time:57560ms step_avg:50.67ms
step:1137/1600 train_time:57649ms step_avg:50.70ms
step:1138/1600 train_time:57733ms step_avg:50.73ms
step:1139/1600 train_time:57827ms step_avg:50.77ms
step:1140/1600 train_time:57910ms step_avg:50.80ms
step:1141/1600 train_time:58001ms step_avg:50.83ms
step:1142/1600 train_time:58083ms step_avg:50.86ms
step:1143/1600 train_time:58170ms step_avg:50.89ms
step:1144/1600 train_time:58255ms step_avg:50.92ms
step:1145/1600 train_time:58344ms step_avg:50.96ms
step:1146/1600 train_time:58429ms step_avg:50.98ms
step:1147/1600 train_time:58517ms step_avg:51.02ms
step:1148/1600 train_time:58602ms step_avg:51.05ms
step:1149/1600 train_time:58690ms step_avg:51.08ms
step:1150/1600 train_time:58775ms step_avg:51.11ms
step:1151/1600 train_time:58863ms step_avg:51.14ms
step:1152/1600 train_time:58950ms step_avg:51.17ms
step:1153/1600 train_time:59037ms step_avg:51.20ms
step:1154/1600 train_time:59126ms step_avg:51.24ms
step:1155/1600 train_time:59211ms step_avg:51.27ms
step:1156/1600 train_time:59297ms step_avg:51.29ms
step:1157/1600 train_time:59385ms step_avg:51.33ms
step:1158/1600 train_time:59470ms step_avg:51.36ms
step:1159/1600 train_time:59558ms step_avg:51.39ms
step:1160/1600 train_time:59643ms step_avg:51.42ms
step:1161/1600 train_time:59731ms step_avg:51.45ms
step:1162/1600 train_time:59817ms step_avg:51.48ms
step:1163/1600 train_time:59904ms step_avg:51.51ms
step:1164/1600 train_time:59989ms step_avg:51.54ms
step:1165/1600 train_time:60077ms step_avg:51.57ms
step:1166/1600 train_time:60163ms step_avg:51.60ms
step:1167/1600 train_time:60251ms step_avg:51.63ms
step:1168/1600 train_time:60337ms step_avg:51.66ms
step:1169/1600 train_time:60425ms step_avg:51.69ms
step:1170/1600 train_time:60510ms step_avg:51.72ms
step:1171/1600 train_time:60598ms step_avg:51.75ms
step:1172/1600 train_time:60683ms step_avg:51.78ms
step:1173/1600 train_time:60770ms step_avg:51.81ms
step:1174/1600 train_time:60855ms step_avg:51.84ms
step:1175/1600 train_time:60944ms step_avg:51.87ms
step:1176/1600 train_time:61029ms step_avg:51.90ms
step:1177/1600 train_time:61117ms step_avg:51.93ms
step:1178/1600 train_time:61202ms step_avg:51.95ms
step:1179/1600 train_time:61290ms step_avg:51.98ms
step:1180/1600 train_time:61375ms step_avg:52.01ms
step:1181/1600 train_time:61464ms step_avg:52.04ms
step:1182/1600 train_time:61549ms step_avg:52.07ms
step:1183/1600 train_time:61638ms step_avg:52.10ms
step:1184/1600 train_time:61722ms step_avg:52.13ms
step:1185/1600 train_time:61810ms step_avg:52.16ms
step:1186/1600 train_time:61895ms step_avg:52.19ms
step:1187/1600 train_time:61984ms step_avg:52.22ms
step:1188/1600 train_time:62069ms step_avg:52.25ms
step:1189/1600 train_time:62157ms step_avg:52.28ms
step:1190/1600 train_time:62242ms step_avg:52.30ms
step:1191/1600 train_time:62331ms step_avg:52.33ms
step:1192/1600 train_time:62415ms step_avg:52.36ms
step:1193/1600 train_time:62503ms step_avg:52.39ms
step:1194/1600 train_time:62588ms step_avg:52.42ms
step:1195/1600 train_time:62676ms step_avg:52.45ms
step:1196/1600 train_time:62762ms step_avg:52.48ms
step:1197/1600 train_time:62849ms step_avg:52.51ms
step:1198/1600 train_time:62934ms step_avg:52.53ms
step:1199/1600 train_time:63022ms step_avg:52.56ms
step:1200/1600 train_time:63107ms step_avg:52.59ms
step:1201/1600 train_time:63194ms step_avg:52.62ms
step:1202/1600 train_time:63280ms step_avg:52.65ms
step:1203/1600 train_time:63367ms step_avg:52.67ms
step:1204/1600 train_time:63452ms step_avg:52.70ms
step:1205/1600 train_time:63541ms step_avg:52.73ms
step:1206/1600 train_time:63625ms step_avg:52.76ms
step:1207/1600 train_time:63714ms step_avg:52.79ms
step:1208/1600 train_time:63799ms step_avg:52.81ms
step:1209/1600 train_time:63887ms step_avg:52.84ms
step:1210/1600 train_time:63972ms step_avg:52.87ms
step:1211/1600 train_time:64060ms step_avg:52.90ms
step:1212/1600 train_time:64145ms step_avg:52.92ms
step:1213/1600 train_time:64233ms step_avg:52.95ms
step:1214/1600 train_time:64318ms step_avg:52.98ms
step:1215/1600 train_time:64406ms step_avg:53.01ms
step:1216/1600 train_time:64491ms step_avg:53.04ms
step:1217/1600 train_time:64579ms step_avg:53.06ms
step:1218/1600 train_time:64664ms step_avg:53.09ms
step:1219/1600 train_time:64753ms step_avg:53.12ms
step:1220/1600 train_time:64840ms step_avg:53.15ms
step:1221/1600 train_time:64936ms step_avg:53.18ms
step:1222/1600 train_time:65012ms step_avg:53.20ms
step:1223/1600 train_time:65100ms step_avg:53.23ms
step:1224/1600 train_time:65185ms step_avg:53.26ms
step:1225/1600 train_time:65274ms step_avg:53.28ms
step:1226/1600 train_time:65360ms step_avg:53.31ms
step:1227/1600 train_time:65448ms step_avg:53.34ms
step:1228/1600 train_time:65533ms step_avg:53.37ms
step:1229/1600 train_time:65623ms step_avg:53.40ms
step:1230/1600 train_time:65708ms step_avg:53.42ms
step:1231/1600 train_time:65795ms step_avg:53.45ms
step:1232/1600 train_time:65880ms step_avg:53.47ms
step:1233/1600 train_time:65967ms step_avg:53.50ms
step:1234/1600 train_time:66053ms step_avg:53.53ms
step:1235/1600 train_time:66141ms step_avg:53.56ms
step:1236/1600 train_time:66226ms step_avg:53.58ms
step:1237/1600 train_time:66314ms step_avg:53.61ms
step:1238/1600 train_time:66399ms step_avg:53.63ms
step:1239/1600 train_time:66487ms step_avg:53.66ms
step:1240/1600 train_time:66572ms step_avg:53.69ms
step:1241/1600 train_time:66661ms step_avg:53.72ms
step:1242/1600 train_time:66746ms step_avg:53.74ms
step:1243/1600 train_time:66835ms step_avg:53.77ms
step:1244/1600 train_time:66918ms step_avg:53.79ms
step:1245/1600 train_time:67006ms step_avg:53.82ms
step:1246/1600 train_time:67091ms step_avg:53.84ms
step:1247/1600 train_time:67179ms step_avg:53.87ms
step:1248/1600 train_time:67264ms step_avg:53.90ms
step:1249/1600 train_time:67352ms step_avg:53.92ms
step:1250/1600 train_time:67438ms step_avg:53.95ms
step:1250/1600 val_loss:3.4162 train_time:67509ms step_avg:54.01ms
step:1251/1600 train_time:67530ms step_avg:53.98ms
step:1252/1600 train_time:67616ms step_avg:54.01ms
step:1253/1600 train_time:67709ms step_avg:54.04ms
step:1254/1600 train_time:67794ms step_avg:54.06ms
step:1255/1600 train_time:67881ms step_avg:54.09ms
step:1256/1600 train_time:67966ms step_avg:54.11ms
step:1257/1600 train_time:68052ms step_avg:54.14ms
step:1258/1600 train_time:68137ms step_avg:54.16ms
step:1259/1600 train_time:68224ms step_avg:54.19ms
step:1260/1600 train_time:68310ms step_avg:54.21ms
step:1261/1600 train_time:68396ms step_avg:54.24ms
step:1262/1600 train_time:68483ms step_avg:54.27ms
step:1263/1600 train_time:68571ms step_avg:54.29ms
step:1264/1600 train_time:68659ms step_avg:54.32ms
step:1265/1600 train_time:68748ms step_avg:54.35ms
step:1266/1600 train_time:68833ms step_avg:54.37ms
step:1267/1600 train_time:68921ms step_avg:54.40ms
step:1268/1600 train_time:69005ms step_avg:54.42ms
step:1269/1600 train_time:69093ms step_avg:54.45ms
step:1270/1600 train_time:69179ms step_avg:54.47ms
step:1271/1600 train_time:69264ms step_avg:54.50ms
step:1272/1600 train_time:69349ms step_avg:54.52ms
step:1273/1600 train_time:69437ms step_avg:54.55ms
step:1274/1600 train_time:69523ms step_avg:54.57ms
step:1275/1600 train_time:69613ms step_avg:54.60ms
step:1276/1600 train_time:69699ms step_avg:54.62ms
step:1277/1600 train_time:69788ms step_avg:54.65ms
step:1278/1600 train_time:69873ms step_avg:54.67ms
step:1279/1600 train_time:69963ms step_avg:54.70ms
step:1280/1600 train_time:70046ms step_avg:54.72ms
step:1281/1600 train_time:70134ms step_avg:54.75ms
step:1282/1600 train_time:70218ms step_avg:54.77ms
step:1283/1600 train_time:70306ms step_avg:54.80ms
step:1284/1600 train_time:70390ms step_avg:54.82ms
step:1285/1600 train_time:70479ms step_avg:54.85ms
step:1286/1600 train_time:70565ms step_avg:54.87ms
step:1287/1600 train_time:70654ms step_avg:54.90ms
step:1288/1600 train_time:70740ms step_avg:54.92ms
step:1289/1600 train_time:70828ms step_avg:54.95ms
step:1290/1600 train_time:70913ms step_avg:54.97ms
step:1291/1600 train_time:71002ms step_avg:55.00ms
step:1292/1600 train_time:71086ms step_avg:55.02ms
step:1293/1600 train_time:71174ms step_avg:55.05ms
step:1294/1600 train_time:71258ms step_avg:55.07ms
step:1295/1600 train_time:71346ms step_avg:55.09ms
step:1296/1600 train_time:71431ms step_avg:55.12ms
step:1297/1600 train_time:71520ms step_avg:55.14ms
step:1298/1600 train_time:71606ms step_avg:55.17ms
step:1299/1600 train_time:71694ms step_avg:55.19ms
step:1300/1600 train_time:71779ms step_avg:55.21ms
step:1301/1600 train_time:71867ms step_avg:55.24ms
step:1302/1600 train_time:71953ms step_avg:55.26ms
step:1303/1600 train_time:72041ms step_avg:55.29ms
step:1304/1600 train_time:72126ms step_avg:55.31ms
step:1305/1600 train_time:72213ms step_avg:55.34ms
step:1306/1600 train_time:72298ms step_avg:55.36ms
step:1307/1600 train_time:72386ms step_avg:55.38ms
step:1308/1600 train_time:72473ms step_avg:55.41ms
step:1309/1600 train_time:72561ms step_avg:55.43ms
step:1310/1600 train_time:72647ms step_avg:55.46ms
step:1311/1600 train_time:72735ms step_avg:55.48ms
step:1312/1600 train_time:72820ms step_avg:55.50ms
step:1313/1600 train_time:72908ms step_avg:55.53ms
step:1314/1600 train_time:72995ms step_avg:55.55ms
step:1315/1600 train_time:73083ms step_avg:55.58ms
step:1316/1600 train_time:73168ms step_avg:55.60ms
step:1317/1600 train_time:73254ms step_avg:55.62ms
step:1318/1600 train_time:73339ms step_avg:55.64ms
step:1319/1600 train_time:73426ms step_avg:55.67ms
step:1320/1600 train_time:73511ms step_avg:55.69ms
step:1321/1600 train_time:73600ms step_avg:55.72ms
step:1322/1600 train_time:73686ms step_avg:55.74ms
step:1323/1600 train_time:73774ms step_avg:55.76ms
step:1324/1600 train_time:73859ms step_avg:55.78ms
step:1325/1600 train_time:73947ms step_avg:55.81ms
step:1326/1600 train_time:74031ms step_avg:55.83ms
step:1327/1600 train_time:74120ms step_avg:55.86ms
step:1328/1600 train_time:74204ms step_avg:55.88ms
step:1329/1600 train_time:74294ms step_avg:55.90ms
step:1330/1600 train_time:74377ms step_avg:55.92ms
step:1331/1600 train_time:74467ms step_avg:55.95ms
step:1332/1600 train_time:74551ms step_avg:55.97ms
step:1333/1600 train_time:74640ms step_avg:55.99ms
step:1334/1600 train_time:74724ms step_avg:56.02ms
step:1335/1600 train_time:74813ms step_avg:56.04ms
step:1336/1600 train_time:74899ms step_avg:56.06ms
step:1337/1600 train_time:74986ms step_avg:56.09ms
step:1338/1600 train_time:75071ms step_avg:56.11ms
step:1339/1600 train_time:75159ms step_avg:56.13ms
step:1340/1600 train_time:75244ms step_avg:56.15ms
step:1341/1600 train_time:75331ms step_avg:56.18ms
step:1342/1600 train_time:75416ms step_avg:56.20ms
step:1343/1600 train_time:75504ms step_avg:56.22ms
step:1344/1600 train_time:75589ms step_avg:56.24ms
step:1345/1600 train_time:75678ms step_avg:56.27ms
step:1346/1600 train_time:75762ms step_avg:56.29ms
step:1347/1600 train_time:75850ms step_avg:56.31ms
step:1348/1600 train_time:75936ms step_avg:56.33ms
step:1349/1600 train_time:76025ms step_avg:56.36ms
step:1350/1600 train_time:76109ms step_avg:56.38ms
step:1351/1600 train_time:76198ms step_avg:56.40ms
step:1352/1600 train_time:76282ms step_avg:56.42ms
step:1353/1600 train_time:76371ms step_avg:56.45ms
step:1354/1600 train_time:76455ms step_avg:56.47ms
step:1355/1600 train_time:76544ms step_avg:56.49ms
step:1356/1600 train_time:76630ms step_avg:56.51ms
step:1357/1600 train_time:76718ms step_avg:56.54ms
step:1358/1600 train_time:76802ms step_avg:56.56ms
step:1359/1600 train_time:76891ms step_avg:56.58ms
step:1360/1600 train_time:76976ms step_avg:56.60ms
step:1361/1600 train_time:77063ms step_avg:56.62ms
step:1362/1600 train_time:77148ms step_avg:56.64ms
step:1363/1600 train_time:77236ms step_avg:56.67ms
step:1364/1600 train_time:77321ms step_avg:56.69ms
step:1365/1600 train_time:77408ms step_avg:56.71ms
step:1366/1600 train_time:77499ms step_avg:56.73ms
step:1367/1600 train_time:77584ms step_avg:56.76ms
step:1368/1600 train_time:77668ms step_avg:56.78ms
step:1369/1600 train_time:77757ms step_avg:56.80ms
step:1370/1600 train_time:77841ms step_avg:56.82ms
step:1371/1600 train_time:77930ms step_avg:56.84ms
step:1372/1600 train_time:78015ms step_avg:56.86ms
step:1373/1600 train_time:78103ms step_avg:56.89ms
step:1374/1600 train_time:78189ms step_avg:56.91ms
step:1375/1600 train_time:78279ms step_avg:56.93ms
step:1376/1600 train_time:78364ms step_avg:56.95ms
step:1377/1600 train_time:78451ms step_avg:56.97ms
step:1378/1600 train_time:78535ms step_avg:56.99ms
step:1379/1600 train_time:78623ms step_avg:57.01ms
step:1380/1600 train_time:78708ms step_avg:57.03ms
step:1381/1600 train_time:78796ms step_avg:57.06ms
step:1382/1600 train_time:78881ms step_avg:57.08ms
step:1383/1600 train_time:78970ms step_avg:57.10ms
step:1384/1600 train_time:79055ms step_avg:57.12ms
step:1385/1600 train_time:79144ms step_avg:57.14ms
step:1386/1600 train_time:79229ms step_avg:57.16ms
step:1387/1600 train_time:79319ms step_avg:57.19ms
step:1388/1600 train_time:79405ms step_avg:57.21ms
step:1389/1600 train_time:79491ms step_avg:57.23ms
step:1390/1600 train_time:79576ms step_avg:57.25ms
step:1391/1600 train_time:79664ms step_avg:57.27ms
step:1392/1600 train_time:79749ms step_avg:57.29ms
step:1393/1600 train_time:79838ms step_avg:57.31ms
step:1394/1600 train_time:79922ms step_avg:57.33ms
step:1395/1600 train_time:80011ms step_avg:57.36ms
step:1396/1600 train_time:80097ms step_avg:57.38ms
step:1397/1600 train_time:80185ms step_avg:57.40ms
step:1398/1600 train_time:80271ms step_avg:57.42ms
step:1399/1600 train_time:80360ms step_avg:57.44ms
step:1400/1600 train_time:80445ms step_avg:57.46ms
step:1401/1600 train_time:80533ms step_avg:57.48ms
step:1402/1600 train_time:80618ms step_avg:57.50ms
step:1403/1600 train_time:80705ms step_avg:57.52ms
step:1404/1600 train_time:80790ms step_avg:57.54ms
step:1405/1600 train_time:80879ms step_avg:57.56ms
step:1406/1600 train_time:80964ms step_avg:57.58ms
step:1407/1600 train_time:81052ms step_avg:57.61ms
step:1408/1600 train_time:81137ms step_avg:57.63ms
step:1409/1600 train_time:81225ms step_avg:57.65ms
step:1410/1600 train_time:81311ms step_avg:57.67ms
step:1411/1600 train_time:81400ms step_avg:57.69ms
step:1412/1600 train_time:81485ms step_avg:57.71ms
step:1413/1600 train_time:81572ms step_avg:57.73ms
step:1414/1600 train_time:81657ms step_avg:57.75ms
step:1415/1600 train_time:81745ms step_avg:57.77ms
step:1416/1600 train_time:81829ms step_avg:57.79ms
step:1417/1600 train_time:81919ms step_avg:57.81ms
step:1418/1600 train_time:82003ms step_avg:57.83ms
step:1419/1600 train_time:82091ms step_avg:57.85ms
step:1420/1600 train_time:82176ms step_avg:57.87ms
step:1421/1600 train_time:82264ms step_avg:57.89ms
step:1422/1600 train_time:82350ms step_avg:57.91ms
step:1423/1600 train_time:82438ms step_avg:57.93ms
step:1424/1600 train_time:82524ms step_avg:57.95ms
step:1425/1600 train_time:82611ms step_avg:57.97ms
step:1426/1600 train_time:82696ms step_avg:57.99ms
step:1427/1600 train_time:82784ms step_avg:58.01ms
step:1428/1600 train_time:82869ms step_avg:58.03ms
step:1429/1600 train_time:82957ms step_avg:58.05ms
step:1430/1600 train_time:83042ms step_avg:58.07ms
step:1431/1600 train_time:83130ms step_avg:58.09ms
step:1432/1600 train_time:83218ms step_avg:58.11ms
step:1433/1600 train_time:83305ms step_avg:58.13ms
step:1434/1600 train_time:83390ms step_avg:58.15ms
step:1435/1600 train_time:83480ms step_avg:58.17ms
step:1436/1600 train_time:83565ms step_avg:58.19ms
step:1437/1600 train_time:83653ms step_avg:58.21ms
step:1438/1600 train_time:83738ms step_avg:58.23ms
step:1439/1600 train_time:83825ms step_avg:58.25ms
step:1440/1600 train_time:83912ms step_avg:58.27ms
step:1441/1600 train_time:84000ms step_avg:58.29ms
step:1442/1600 train_time:84084ms step_avg:58.31ms
step:1443/1600 train_time:84173ms step_avg:58.33ms
step:1444/1600 train_time:84259ms step_avg:58.35ms
step:1445/1600 train_time:84347ms step_avg:58.37ms
step:1446/1600 train_time:84433ms step_avg:58.39ms
step:1447/1600 train_time:84523ms step_avg:58.41ms
step:1448/1600 train_time:84607ms step_avg:58.43ms
step:1449/1600 train_time:84695ms step_avg:58.45ms
step:1450/1600 train_time:84780ms step_avg:58.47ms
step:1451/1600 train_time:84868ms step_avg:58.49ms
step:1452/1600 train_time:84954ms step_avg:58.51ms
step:1453/1600 train_time:85042ms step_avg:58.53ms
step:1454/1600 train_time:85127ms step_avg:58.55ms
step:1455/1600 train_time:85215ms step_avg:58.57ms
step:1456/1600 train_time:85301ms step_avg:58.59ms
step:1457/1600 train_time:85388ms step_avg:58.61ms
step:1458/1600 train_time:85476ms step_avg:58.63ms
step:1459/1600 train_time:85563ms step_avg:58.64ms
step:1460/1600 train_time:85648ms step_avg:58.66ms
step:1461/1600 train_time:85736ms step_avg:58.68ms
step:1462/1600 train_time:85821ms step_avg:58.70ms
step:1463/1600 train_time:85909ms step_avg:58.72ms
step:1464/1600 train_time:85996ms step_avg:58.74ms
step:1465/1600 train_time:86083ms step_avg:58.76ms
step:1466/1600 train_time:86168ms step_avg:58.78ms
step:1467/1600 train_time:86257ms step_avg:58.80ms
step:1468/1600 train_time:86345ms step_avg:58.82ms
step:1469/1600 train_time:86431ms step_avg:58.84ms
step:1470/1600 train_time:86516ms step_avg:58.85ms
step:1471/1600 train_time:86606ms step_avg:58.88ms
step:1472/1600 train_time:86690ms step_avg:58.89ms
step:1473/1600 train_time:86779ms step_avg:58.91ms
step:1474/1600 train_time:86864ms step_avg:58.93ms
step:1475/1600 train_time:86952ms step_avg:58.95ms
step:1476/1600 train_time:87037ms step_avg:58.97ms
step:1477/1600 train_time:87126ms step_avg:58.99ms
step:1478/1600 train_time:87211ms step_avg:59.01ms
step:1479/1600 train_time:87300ms step_avg:59.03ms
step:1480/1600 train_time:87386ms step_avg:59.04ms
step:1481/1600 train_time:87475ms step_avg:59.06ms
step:1482/1600 train_time:87559ms step_avg:59.08ms
step:1483/1600 train_time:87647ms step_avg:59.10ms
step:1484/1600 train_time:87733ms step_avg:59.12ms
step:1485/1600 train_time:87821ms step_avg:59.14ms
step:1486/1600 train_time:87905ms step_avg:59.16ms
step:1487/1600 train_time:87994ms step_avg:59.18ms
step:1488/1600 train_time:88083ms step_avg:59.20ms
step:1489/1600 train_time:88167ms step_avg:59.21ms
step:1490/1600 train_time:88252ms step_avg:59.23ms
step:1491/1600 train_time:88341ms step_avg:59.25ms
step:1492/1600 train_time:88426ms step_avg:59.27ms
step:1493/1600 train_time:88514ms step_avg:59.29ms
step:1494/1600 train_time:88598ms step_avg:59.30ms
step:1495/1600 train_time:88686ms step_avg:59.32ms
step:1496/1600 train_time:88772ms step_avg:59.34ms
step:1497/1600 train_time:88860ms step_avg:59.36ms
step:1498/1600 train_time:88945ms step_avg:59.38ms
step:1499/1600 train_time:89033ms step_avg:59.39ms
step:1500/1600 train_time:89117ms step_avg:59.41ms
step:1500/1600 val_loss:3.3065 train_time:89189ms step_avg:59.46ms
step:1501/1600 train_time:89210ms step_avg:59.43ms
step:1502/1600 train_time:89297ms step_avg:59.45ms
step:1503/1600 train_time:89390ms step_avg:59.47ms
step:1504/1600 train_time:89476ms step_avg:59.49ms
step:1505/1600 train_time:89564ms step_avg:59.51ms
step:1506/1600 train_time:89648ms step_avg:59.53ms
step:1507/1600 train_time:89735ms step_avg:59.55ms
step:1508/1600 train_time:89820ms step_avg:59.56ms
step:1509/1600 train_time:89906ms step_avg:59.58ms
step:1510/1600 train_time:89990ms step_avg:59.60ms
step:1511/1600 train_time:90078ms step_avg:59.61ms
step:1512/1600 train_time:90164ms step_avg:59.63ms
step:1513/1600 train_time:90255ms step_avg:59.65ms
step:1514/1600 train_time:90343ms step_avg:59.67ms
step:1515/1600 train_time:90432ms step_avg:59.69ms
step:1516/1600 train_time:90518ms step_avg:59.71ms
step:1517/1600 train_time:90607ms step_avg:59.73ms
step:1518/1600 train_time:90690ms step_avg:59.74ms
step:1519/1600 train_time:90778ms step_avg:59.76ms
step:1520/1600 train_time:90862ms step_avg:59.78ms
step:1521/1600 train_time:90948ms step_avg:59.80ms
step:1522/1600 train_time:91032ms step_avg:59.81ms
step:1523/1600 train_time:91121ms step_avg:59.83ms
step:1524/1600 train_time:91206ms step_avg:59.85ms
step:1525/1600 train_time:91295ms step_avg:59.87ms
step:1526/1600 train_time:91381ms step_avg:59.88ms
step:1527/1600 train_time:91469ms step_avg:59.90ms
step:1528/1600 train_time:91555ms step_avg:59.92ms
step:1529/1600 train_time:91643ms step_avg:59.94ms
step:1530/1600 train_time:91728ms step_avg:59.95ms
step:1531/1600 train_time:91815ms step_avg:59.97ms
step:1532/1600 train_time:91900ms step_avg:59.99ms
step:1533/1600 train_time:91988ms step_avg:60.01ms
step:1534/1600 train_time:92074ms step_avg:60.02ms
step:1535/1600 train_time:92163ms step_avg:60.04ms
step:1536/1600 train_time:92248ms step_avg:60.06ms
step:1537/1600 train_time:92337ms step_avg:60.08ms
step:1538/1600 train_time:92423ms step_avg:60.09ms
step:1539/1600 train_time:92512ms step_avg:60.11ms
step:1540/1600 train_time:92597ms step_avg:60.13ms
step:1541/1600 train_time:92685ms step_avg:60.15ms
step:1542/1600 train_time:92770ms step_avg:60.16ms
step:1543/1600 train_time:92858ms step_avg:60.18ms
step:1544/1600 train_time:92942ms step_avg:60.20ms
step:1545/1600 train_time:93030ms step_avg:60.21ms
step:1546/1600 train_time:93116ms step_avg:60.23ms
step:1547/1600 train_time:93204ms step_avg:60.25ms
step:1548/1600 train_time:93290ms step_avg:60.26ms
step:1549/1600 train_time:93379ms step_avg:60.28ms
step:1550/1600 train_time:93464ms step_avg:60.30ms
step:1551/1600 train_time:93553ms step_avg:60.32ms
step:1552/1600 train_time:93638ms step_avg:60.33ms
step:1553/1600 train_time:93726ms step_avg:60.35ms
step:1554/1600 train_time:93813ms step_avg:60.37ms
step:1555/1600 train_time:93902ms step_avg:60.39ms
step:1556/1600 train_time:93986ms step_avg:60.40ms
step:1557/1600 train_time:94073ms step_avg:60.42ms
step:1558/1600 train_time:94158ms step_avg:60.44ms
step:1559/1600 train_time:94248ms step_avg:60.45ms
step:1560/1600 train_time:94332ms step_avg:60.47ms
step:1561/1600 train_time:94430ms step_avg:60.49ms
step:1562/1600 train_time:94512ms step_avg:60.51ms
step:1563/1600 train_time:94600ms step_avg:60.52ms
step:1564/1600 train_time:94686ms step_avg:60.54ms
step:1565/1600 train_time:94773ms step_avg:60.56ms
step:1566/1600 train_time:94859ms step_avg:60.57ms
step:1567/1600 train_time:94947ms step_avg:60.59ms
step:1568/1600 train_time:95033ms step_avg:60.61ms
step:1569/1600 train_time:95122ms step_avg:60.63ms
step:1570/1600 train_time:95206ms step_avg:60.64ms
step:1571/1600 train_time:95296ms step_avg:60.66ms
step:1572/1600 train_time:95383ms step_avg:60.68ms
step:1573/1600 train_time:95471ms step_avg:60.69ms
step:1574/1600 train_time:95558ms step_avg:60.71ms
step:1575/1600 train_time:95646ms step_avg:60.73ms
step:1576/1600 train_time:95732ms step_avg:60.74ms
step:1577/1600 train_time:95823ms step_avg:60.76ms
step:1578/1600 train_time:95907ms step_avg:60.78ms
step:1579/1600 train_time:95995ms step_avg:60.79ms
step:1580/1600 train_time:96080ms step_avg:60.81ms
step:1581/1600 train_time:96168ms step_avg:60.83ms
step:1582/1600 train_time:96254ms step_avg:60.84ms
step:1583/1600 train_time:96344ms step_avg:60.86ms
step:1584/1600 train_time:96430ms step_avg:60.88ms
step:1585/1600 train_time:96518ms step_avg:60.89ms
step:1586/1600 train_time:96604ms step_avg:60.91ms
step:1587/1600 train_time:96692ms step_avg:60.93ms
step:1588/1600 train_time:96780ms step_avg:60.94ms
step:1589/1600 train_time:96867ms step_avg:60.96ms
step:1590/1600 train_time:96953ms step_avg:60.98ms
step:1591/1600 train_time:97041ms step_avg:60.99ms
step:1592/1600 train_time:97129ms step_avg:61.01ms
step:1593/1600 train_time:97215ms step_avg:61.03ms
step:1594/1600 train_time:97304ms step_avg:61.04ms
step:1595/1600 train_time:97391ms step_avg:61.06ms
step:1596/1600 train_time:97475ms step_avg:61.07ms
step:1597/1600 train_time:97564ms step_avg:61.09ms
step:1598/1600 train_time:97649ms step_avg:61.11ms
step:1599/1600 train_time:97738ms step_avg:61.12ms
step:1600/1600 train_time:97824ms step_avg:61.14ms
step:1600/1600 val_loss:3.2770 train_time:97895ms step_avg:61.18ms
peak memory allocated: 30789 MiB reserved: 46138 MiB