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:26:38 2026       
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 570.148.08             Driver Version: 570.148.08     CUDA Version: 12.8     |
|-----------------------------------------+------------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  NVIDIA H100 80GB HBM3          On  |   00000000:61:00.0 Off |                    0 |
| N/A   34C    P0            120W /  700W |    1519MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   1  NVIDIA H100 80GB HBM3          On  |   00000000:62:00.0 Off |                    0 |
| N/A   38C    P0            121W /  700W |    1519MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   2  NVIDIA H100 80GB HBM3          On  |   00000000:63:00.0 Off |                    0 |
| N/A   40C    P0            120W /  700W |    1519MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   3  NVIDIA H100 80GB HBM3          On  |   00000000:64:00.0 Off |                    0 |
| N/A   34C    P0            120W /  700W |    1519MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   4  NVIDIA H100 80GB HBM3          On  |   00000000:6A:00.0 Off |                    0 |
| N/A   34C    P0            119W /  700W |    1519MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   5  NVIDIA H100 80GB HBM3          On  |   00000000:6B:00.0 Off |                    0 |
| N/A   41C    P0            132W /  700W |    1519MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   6  NVIDIA H100 80GB HBM3          On  |   00000000:6C:00.0 Off |                    0 |
| N/A   39C    P0            123W /  700W |    1519MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   7  NVIDIA H100 80GB HBM3          On  |   00000000:6D:00.0 Off |                    0 |
| N/A   36C    P0            118W /  700W |    1519MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
                                                                                         
+-----------------------------------------------------------------------------------------+
| Processes:                                                                              |
|  GPU   GI   CI              PID   Type   Process name                        GPU Memory |
|        ID   ID                                                               Usage      |
|=========================================================================================|
|    0   N/A  N/A          307242      C   /usr/bin/python3                       1510MiB |
|    1   N/A  N/A          307243      C   /usr/bin/python3                       1510MiB |
|    2   N/A  N/A          307244      C   /usr/bin/python3                       1510MiB |
|    3   N/A  N/A          307245      C   /usr/bin/python3                       1510MiB |
|    4   N/A  N/A          307246      C   /usr/bin/python3                       1510MiB |
|    5   N/A  N/A          307247      C   /usr/bin/python3                       1510MiB |
|    6   N/A  N/A          307248      C   /usr/bin/python3                       1510MiB |
|    7   N/A  N/A          307249      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.8298 train_time:0ms step_avg:0.04ms
step:1/1600 train_time:85ms step_avg:84.73ms
step:2/1600 train_time:105ms step_avg:52.74ms
step:3/1600 train_time:124ms step_avg:41.42ms
step:4/1600 train_time:156ms step_avg:38.91ms
step:5/1600 train_time:185ms step_avg:37.06ms
step:6/1600 train_time:308ms step_avg:51.34ms
step:7/1600 train_time:325ms step_avg:46.38ms
step:8/1600 train_time:449ms step_avg:56.09ms
step:9/1600 train_time:479ms step_avg:53.23ms
step:10/1600 train_time:516ms step_avg:51.62ms
step:11/1600 train_time:547ms step_avg:49.76ms
step:12/1600 train_time:584ms step_avg:48.67ms
step:13/1600 train_time:615ms step_avg:47.32ms
step:14/1600 train_time:653ms step_avg:46.62ms
step:15/1600 train_time:684ms step_avg:45.61ms
step:16/1600 train_time:721ms step_avg:45.07ms
step:17/1600 train_time:752ms step_avg:44.24ms
step:18/1600 train_time:789ms step_avg:43.84ms
step:19/1600 train_time:820ms step_avg:43.15ms
step:20/1600 train_time:857ms step_avg:42.87ms
step:21/1600 train_time:889ms step_avg:42.32ms
step:22/1600 train_time:926ms step_avg:42.07ms
step:23/1600 train_time:957ms step_avg:41.59ms
step:24/1600 train_time:993ms step_avg:41.39ms
step:25/1600 train_time:1024ms step_avg:40.97ms
step:26/1600 train_time:1061ms step_avg:40.82ms
step:27/1600 train_time:1092ms step_avg:40.46ms
step:28/1600 train_time:1130ms step_avg:40.34ms
step:29/1600 train_time:1160ms step_avg:40.02ms
step:30/1600 train_time:1197ms step_avg:39.91ms
step:31/1600 train_time:1229ms step_avg:39.63ms
step:32/1600 train_time:1265ms step_avg:39.54ms
step:33/1600 train_time:1296ms step_avg:39.28ms
step:34/1600 train_time:1334ms step_avg:39.22ms
step:35/1600 train_time:1365ms step_avg:38.99ms
step:36/1600 train_time:1401ms step_avg:38.93ms
step:37/1600 train_time:1433ms step_avg:38.73ms
step:38/1600 train_time:1470ms step_avg:38.69ms
step:39/1600 train_time:1501ms step_avg:38.48ms
step:40/1600 train_time:1539ms step_avg:38.47ms
step:41/1600 train_time:1570ms step_avg:38.29ms
step:42/1600 train_time:1607ms step_avg:38.26ms
step:43/1600 train_time:1638ms step_avg:38.10ms
step:44/1600 train_time:1676ms step_avg:38.09ms
step:45/1600 train_time:1707ms step_avg:37.93ms
step:46/1600 train_time:1744ms step_avg:37.91ms
step:47/1600 train_time:1775ms step_avg:37.76ms
step:48/1600 train_time:1812ms step_avg:37.76ms
step:49/1600 train_time:1843ms step_avg:37.62ms
step:50/1600 train_time:1880ms step_avg:37.61ms
step:51/1600 train_time:1911ms step_avg:37.47ms
step:52/1600 train_time:1948ms step_avg:37.46ms
step:53/1600 train_time:1979ms step_avg:37.33ms
step:54/1600 train_time:2016ms step_avg:37.33ms
step:55/1600 train_time:2047ms step_avg:37.21ms
step:56/1600 train_time:2084ms step_avg:37.21ms
step:57/1600 train_time:2115ms step_avg:37.10ms
step:58/1600 train_time:2152ms step_avg:37.11ms
step:59/1600 train_time:2183ms step_avg:37.01ms
step:60/1600 train_time:2220ms step_avg:37.00ms
step:61/1600 train_time:2251ms step_avg:36.90ms
step:62/1600 train_time:2288ms step_avg:36.91ms
step:63/1600 train_time:2319ms step_avg:36.81ms
step:64/1600 train_time:2356ms step_avg:36.81ms
step:65/1600 train_time:2387ms step_avg:36.73ms
step:66/1600 train_time:2424ms step_avg:36.73ms
step:67/1600 train_time:2455ms step_avg:36.65ms
step:68/1600 train_time:2492ms step_avg:36.65ms
step:69/1600 train_time:2523ms step_avg:36.56ms
step:70/1600 train_time:2560ms step_avg:36.57ms
step:71/1600 train_time:2591ms step_avg:36.49ms
step:72/1600 train_time:2628ms step_avg:36.50ms
step:73/1600 train_time:2659ms step_avg:36.43ms
step:74/1600 train_time:2697ms step_avg:36.45ms
step:75/1600 train_time:2728ms step_avg:36.37ms
step:76/1600 train_time:2765ms step_avg:36.38ms
step:77/1600 train_time:2795ms step_avg:36.30ms
step:78/1600 train_time:2832ms step_avg:36.31ms
step:79/1600 train_time:2863ms step_avg:36.24ms
step:80/1600 train_time:2900ms step_avg:36.25ms
step:81/1600 train_time:2931ms step_avg:36.18ms
step:82/1600 train_time:2968ms step_avg:36.19ms
step:83/1600 train_time:2999ms step_avg:36.13ms
step:84/1600 train_time:3037ms step_avg:36.15ms
step:85/1600 train_time:3068ms step_avg:36.09ms
step:86/1600 train_time:3105ms step_avg:36.10ms
step:87/1600 train_time:3136ms step_avg:36.05ms
step:88/1600 train_time:3173ms step_avg:36.06ms
step:89/1600 train_time:3204ms step_avg:36.00ms
step:90/1600 train_time:3242ms step_avg:36.02ms
step:91/1600 train_time:3272ms step_avg:35.96ms
step:92/1600 train_time:3309ms step_avg:35.97ms
step:93/1600 train_time:3340ms step_avg:35.92ms
step:94/1600 train_time:3377ms step_avg:35.93ms
step:95/1600 train_time:3408ms step_avg:35.88ms
step:96/1600 train_time:3445ms step_avg:35.89ms
step:97/1600 train_time:3476ms step_avg:35.83ms
step:98/1600 train_time:3513ms step_avg:35.85ms
step:99/1600 train_time:3544ms step_avg:35.80ms
step:100/1600 train_time:3581ms step_avg:35.81ms
step:101/1600 train_time:3612ms step_avg:35.77ms
step:102/1600 train_time:3649ms step_avg:35.78ms
step:103/1600 train_time:3680ms step_avg:35.73ms
step:104/1600 train_time:3717ms step_avg:35.74ms
step:105/1600 train_time:3748ms step_avg:35.70ms
step:106/1600 train_time:3785ms step_avg:35.71ms
step:107/1600 train_time:3816ms step_avg:35.67ms
step:108/1600 train_time:3853ms step_avg:35.68ms
step:109/1600 train_time:3884ms step_avg:35.63ms
step:110/1600 train_time:3921ms step_avg:35.64ms
step:111/1600 train_time:3952ms step_avg:35.60ms
step:112/1600 train_time:3989ms step_avg:35.62ms
step:113/1600 train_time:4020ms step_avg:35.58ms
step:114/1600 train_time:4057ms step_avg:35.59ms
step:115/1600 train_time:4088ms step_avg:35.55ms
step:116/1600 train_time:4125ms step_avg:35.56ms
step:117/1600 train_time:4156ms step_avg:35.52ms
step:118/1600 train_time:4193ms step_avg:35.54ms
step:119/1600 train_time:4225ms step_avg:35.50ms
step:120/1600 train_time:4262ms step_avg:35.51ms
step:121/1600 train_time:4293ms step_avg:35.48ms
step:122/1600 train_time:4330ms step_avg:35.49ms
step:123/1600 train_time:4361ms step_avg:35.45ms
step:124/1600 train_time:4398ms step_avg:35.47ms
step:125/1600 train_time:4429ms step_avg:35.43ms
step:126/1600 train_time:4466ms step_avg:35.45ms
step:127/1600 train_time:4497ms step_avg:35.41ms
step:128/1600 train_time:4534ms step_avg:35.42ms
step:129/1600 train_time:4565ms step_avg:35.39ms
step:130/1600 train_time:4602ms step_avg:35.40ms
step:131/1600 train_time:4633ms step_avg:35.36ms
step:132/1600 train_time:4670ms step_avg:35.38ms
step:133/1600 train_time:4701ms step_avg:35.35ms
step:134/1600 train_time:4739ms step_avg:35.36ms
step:135/1600 train_time:4769ms step_avg:35.33ms
step:136/1600 train_time:4806ms step_avg:35.34ms
step:137/1600 train_time:4837ms step_avg:35.30ms
step:138/1600 train_time:4874ms step_avg:35.32ms
step:139/1600 train_time:4905ms step_avg:35.29ms
step:140/1600 train_time:4942ms step_avg:35.30ms
step:141/1600 train_time:4973ms step_avg:35.27ms
step:142/1600 train_time:5011ms step_avg:35.29ms
step:143/1600 train_time:5042ms step_avg:35.26ms
step:144/1600 train_time:5079ms step_avg:35.27ms
step:145/1600 train_time:5110ms step_avg:35.24ms
step:146/1600 train_time:5147ms step_avg:35.25ms
step:147/1600 train_time:5177ms step_avg:35.22ms
step:148/1600 train_time:5215ms step_avg:35.23ms
step:149/1600 train_time:5246ms step_avg:35.21ms
step:150/1600 train_time:5284ms step_avg:35.22ms
step:151/1600 train_time:5314ms step_avg:35.19ms
step:152/1600 train_time:5352ms step_avg:35.21ms
step:153/1600 train_time:5382ms step_avg:35.17ms
step:154/1600 train_time:5419ms step_avg:35.19ms
step:155/1600 train_time:5450ms step_avg:35.16ms
step:156/1600 train_time:5487ms step_avg:35.17ms
step:157/1600 train_time:5517ms step_avg:35.14ms
step:158/1600 train_time:5556ms step_avg:35.16ms
step:159/1600 train_time:5586ms step_avg:35.13ms
step:160/1600 train_time:5623ms step_avg:35.15ms
step:161/1600 train_time:5654ms step_avg:35.12ms
step:162/1600 train_time:5691ms step_avg:35.13ms
step:163/1600 train_time:5721ms step_avg:35.10ms
step:164/1600 train_time:5759ms step_avg:35.11ms
step:165/1600 train_time:5789ms step_avg:35.09ms
step:166/1600 train_time:5826ms step_avg:35.10ms
step:167/1600 train_time:5857ms step_avg:35.07ms
step:168/1600 train_time:5894ms step_avg:35.08ms
step:169/1600 train_time:5925ms step_avg:35.06ms
step:170/1600 train_time:5962ms step_avg:35.07ms
step:171/1600 train_time:5993ms step_avg:35.04ms
step:172/1600 train_time:6030ms step_avg:35.06ms
step:173/1600 train_time:6060ms step_avg:35.03ms
step:174/1600 train_time:6097ms step_avg:35.04ms
step:175/1600 train_time:6128ms step_avg:35.02ms
step:176/1600 train_time:6165ms step_avg:35.03ms
step:177/1600 train_time:6196ms step_avg:35.01ms
step:178/1600 train_time:6234ms step_avg:35.02ms
step:179/1600 train_time:6264ms step_avg:35.00ms
step:180/1600 train_time:6301ms step_avg:35.01ms
step:181/1600 train_time:6332ms step_avg:34.98ms
step:182/1600 train_time:6369ms step_avg:35.00ms
step:183/1600 train_time:6400ms step_avg:34.97ms
step:184/1600 train_time:6437ms step_avg:34.98ms
step:185/1600 train_time:6468ms step_avg:34.96ms
step:186/1600 train_time:6505ms step_avg:34.97ms
step:187/1600 train_time:6536ms step_avg:34.95ms
step:188/1600 train_time:6572ms step_avg:34.96ms
step:189/1600 train_time:6603ms step_avg:34.94ms
step:190/1600 train_time:6640ms step_avg:34.95ms
step:191/1600 train_time:6671ms step_avg:34.93ms
step:192/1600 train_time:6708ms step_avg:34.94ms
step:193/1600 train_time:6739ms step_avg:34.92ms
step:194/1600 train_time:6776ms step_avg:34.93ms
step:195/1600 train_time:6807ms step_avg:34.91ms
step:196/1600 train_time:6844ms step_avg:34.92ms
step:197/1600 train_time:6875ms step_avg:34.90ms
step:198/1600 train_time:6912ms step_avg:34.91ms
step:199/1600 train_time:6943ms step_avg:34.89ms
step:200/1600 train_time:6980ms step_avg:34.90ms
step:201/1600 train_time:7010ms step_avg:34.88ms
step:202/1600 train_time:7047ms step_avg:34.89ms
step:203/1600 train_time:7078ms step_avg:34.87ms
step:204/1600 train_time:7115ms step_avg:34.88ms
step:205/1600 train_time:7146ms step_avg:34.86ms
step:206/1600 train_time:7184ms step_avg:34.87ms
step:207/1600 train_time:7216ms step_avg:34.86ms
step:208/1600 train_time:7253ms step_avg:34.87ms
step:209/1600 train_time:7284ms step_avg:34.85ms
step:210/1600 train_time:7320ms step_avg:34.86ms
step:211/1600 train_time:7351ms step_avg:34.84ms
step:212/1600 train_time:7388ms step_avg:34.85ms
step:213/1600 train_time:7419ms step_avg:34.83ms
step:214/1600 train_time:7456ms step_avg:34.84ms
step:215/1600 train_time:7487ms step_avg:34.82ms
step:216/1600 train_time:7524ms step_avg:34.83ms
step:217/1600 train_time:7555ms step_avg:34.82ms
step:218/1600 train_time:7592ms step_avg:34.83ms
step:219/1600 train_time:7623ms step_avg:34.81ms
step:220/1600 train_time:7660ms step_avg:34.82ms
step:221/1600 train_time:7690ms step_avg:34.80ms
step:222/1600 train_time:7727ms step_avg:34.81ms
step:223/1600 train_time:7758ms step_avg:34.79ms
step:224/1600 train_time:7795ms step_avg:34.80ms
step:225/1600 train_time:7826ms step_avg:34.78ms
step:226/1600 train_time:7864ms step_avg:34.79ms
step:227/1600 train_time:7894ms step_avg:34.78ms
step:228/1600 train_time:7931ms step_avg:34.79ms
step:229/1600 train_time:7962ms step_avg:34.77ms
step:230/1600 train_time:7999ms step_avg:34.78ms
step:231/1600 train_time:8029ms step_avg:34.76ms
step:232/1600 train_time:8066ms step_avg:34.77ms
step:233/1600 train_time:8097ms step_avg:34.75ms
step:234/1600 train_time:8134ms step_avg:34.76ms
step:235/1600 train_time:8165ms step_avg:34.74ms
step:236/1600 train_time:8202ms step_avg:34.75ms
step:237/1600 train_time:8233ms step_avg:34.74ms
step:238/1600 train_time:8270ms step_avg:34.75ms
step:239/1600 train_time:8300ms step_avg:34.73ms
step:240/1600 train_time:8337ms step_avg:34.74ms
step:241/1600 train_time:8368ms step_avg:34.72ms
step:242/1600 train_time:8405ms step_avg:34.73ms
step:243/1600 train_time:8436ms step_avg:34.71ms
step:244/1600 train_time:8472ms step_avg:34.72ms
step:245/1600 train_time:8503ms step_avg:34.71ms
step:246/1600 train_time:8540ms step_avg:34.72ms
step:247/1600 train_time:8571ms step_avg:34.70ms
step:248/1600 train_time:8608ms step_avg:34.71ms
step:249/1600 train_time:8639ms step_avg:34.69ms
step:250/1600 train_time:8676ms step_avg:34.71ms
step:250/1600 val_loss:4.5759 train_time:8724ms step_avg:34.90ms
step:251/1600 train_time:8743ms step_avg:34.83ms
step:252/1600 train_time:8761ms step_avg:34.77ms
step:253/1600 train_time:8777ms step_avg:34.69ms
step:254/1600 train_time:8814ms step_avg:34.70ms
step:255/1600 train_time:8846ms step_avg:34.69ms
step:256/1600 train_time:8886ms step_avg:34.71ms
step:257/1600 train_time:8917ms step_avg:34.70ms
step:258/1600 train_time:8954ms step_avg:34.71ms
step:259/1600 train_time:8985ms step_avg:34.69ms
step:260/1600 train_time:9023ms step_avg:34.70ms
step:261/1600 train_time:9053ms step_avg:34.69ms
step:262/1600 train_time:9090ms step_avg:34.69ms
step:263/1600 train_time:9120ms step_avg:34.68ms
step:264/1600 train_time:9158ms step_avg:34.69ms
step:265/1600 train_time:9189ms step_avg:34.67ms
step:266/1600 train_time:9226ms step_avg:34.68ms
step:267/1600 train_time:9257ms step_avg:34.67ms
step:268/1600 train_time:9294ms step_avg:34.68ms
step:269/1600 train_time:9324ms step_avg:34.66ms
step:270/1600 train_time:9361ms step_avg:34.67ms
step:271/1600 train_time:9392ms step_avg:34.66ms
step:272/1600 train_time:9429ms step_avg:34.67ms
step:273/1600 train_time:9460ms step_avg:34.65ms
step:274/1600 train_time:9496ms step_avg:34.66ms
step:275/1600 train_time:9527ms step_avg:34.64ms
step:276/1600 train_time:9564ms step_avg:34.65ms
step:277/1600 train_time:9595ms step_avg:34.64ms
step:278/1600 train_time:9632ms step_avg:34.65ms
step:279/1600 train_time:9662ms step_avg:34.63ms
step:280/1600 train_time:9699ms step_avg:34.64ms
step:281/1600 train_time:9730ms step_avg:34.63ms
step:282/1600 train_time:9767ms step_avg:34.63ms
step:283/1600 train_time:9798ms step_avg:34.62ms
step:284/1600 train_time:9835ms step_avg:34.63ms
step:285/1600 train_time:9865ms step_avg:34.62ms
step:286/1600 train_time:9902ms step_avg:34.62ms
step:287/1600 train_time:9934ms step_avg:34.61ms
step:288/1600 train_time:9971ms step_avg:34.62ms
step:289/1600 train_time:10002ms step_avg:34.61ms
step:290/1600 train_time:10039ms step_avg:34.62ms
step:291/1600 train_time:10069ms step_avg:34.60ms
step:292/1600 train_time:10106ms step_avg:34.61ms
step:293/1600 train_time:10137ms step_avg:34.60ms
step:294/1600 train_time:10174ms step_avg:34.60ms
step:295/1600 train_time:10205ms step_avg:34.59ms
step:296/1600 train_time:10242ms step_avg:34.60ms
step:297/1600 train_time:10273ms step_avg:34.59ms
step:298/1600 train_time:10310ms step_avg:34.60ms
step:299/1600 train_time:10340ms step_avg:34.58ms
step:300/1600 train_time:10377ms step_avg:34.59ms
step:301/1600 train_time:10408ms step_avg:34.58ms
step:302/1600 train_time:10445ms step_avg:34.59ms
step:303/1600 train_time:10476ms step_avg:34.57ms
step:304/1600 train_time:10513ms step_avg:34.58ms
step:305/1600 train_time:10544ms step_avg:34.57ms
step:306/1600 train_time:10581ms step_avg:34.58ms
step:307/1600 train_time:10612ms step_avg:34.57ms
step:308/1600 train_time:10649ms step_avg:34.57ms
step:309/1600 train_time:10680ms step_avg:34.56ms
step:310/1600 train_time:10717ms step_avg:34.57ms
step:311/1600 train_time:10748ms step_avg:34.56ms
step:312/1600 train_time:10785ms step_avg:34.57ms
step:313/1600 train_time:10816ms step_avg:34.56ms
step:314/1600 train_time:10853ms step_avg:34.56ms
step:315/1600 train_time:10884ms step_avg:34.55ms
step:316/1600 train_time:10921ms step_avg:34.56ms
step:317/1600 train_time:10952ms step_avg:34.55ms
step:318/1600 train_time:10989ms step_avg:34.56ms
step:319/1600 train_time:11020ms step_avg:34.55ms
step:320/1600 train_time:11057ms step_avg:34.55ms
step:321/1600 train_time:11088ms step_avg:34.54ms
step:322/1600 train_time:11125ms step_avg:34.55ms
step:323/1600 train_time:11156ms step_avg:34.54ms
step:324/1600 train_time:11193ms step_avg:34.54ms
step:325/1600 train_time:11223ms step_avg:34.53ms
step:326/1600 train_time:11261ms step_avg:34.54ms
step:327/1600 train_time:11292ms step_avg:34.53ms
step:328/1600 train_time:11329ms step_avg:34.54ms
step:329/1600 train_time:11360ms step_avg:34.53ms
step:330/1600 train_time:11398ms step_avg:34.54ms
step:331/1600 train_time:11429ms step_avg:34.53ms
step:332/1600 train_time:11466ms step_avg:34.54ms
step:333/1600 train_time:11497ms step_avg:34.52ms
step:334/1600 train_time:11534ms step_avg:34.53ms
step:335/1600 train_time:11564ms step_avg:34.52ms
step:336/1600 train_time:11601ms step_avg:34.53ms
step:337/1600 train_time:11632ms step_avg:34.52ms
step:338/1600 train_time:11668ms step_avg:34.52ms
step:339/1600 train_time:11699ms step_avg:34.51ms
step:340/1600 train_time:11736ms step_avg:34.52ms
step:341/1600 train_time:11766ms step_avg:34.51ms
step:342/1600 train_time:11804ms step_avg:34.51ms
step:343/1600 train_time:11835ms step_avg:34.50ms
step:344/1600 train_time:11872ms step_avg:34.51ms
step:345/1600 train_time:11902ms step_avg:34.50ms
step:346/1600 train_time:11940ms step_avg:34.51ms
step:347/1600 train_time:11971ms step_avg:34.50ms
step:348/1600 train_time:12008ms step_avg:34.51ms
step:349/1600 train_time:12039ms step_avg:34.49ms
step:350/1600 train_time:12075ms step_avg:34.50ms
step:351/1600 train_time:12106ms step_avg:34.49ms
step:352/1600 train_time:12143ms step_avg:34.50ms
step:353/1600 train_time:12174ms step_avg:34.49ms
step:354/1600 train_time:12211ms step_avg:34.49ms
step:355/1600 train_time:12242ms step_avg:34.48ms
step:356/1600 train_time:12279ms step_avg:34.49ms
step:357/1600 train_time:12310ms step_avg:34.48ms
step:358/1600 train_time:12347ms step_avg:34.49ms
step:359/1600 train_time:12378ms step_avg:34.48ms
step:360/1600 train_time:12415ms step_avg:34.48ms
step:361/1600 train_time:12445ms step_avg:34.47ms
step:362/1600 train_time:12482ms step_avg:34.48ms
step:363/1600 train_time:12513ms step_avg:34.47ms
step:364/1600 train_time:12550ms step_avg:34.48ms
step:365/1600 train_time:12581ms step_avg:34.47ms
step:366/1600 train_time:12618ms step_avg:34.48ms
step:367/1600 train_time:12649ms step_avg:34.47ms
step:368/1600 train_time:12686ms step_avg:34.47ms
step:369/1600 train_time:12717ms step_avg:34.46ms
step:370/1600 train_time:12754ms step_avg:34.47ms
step:371/1600 train_time:12785ms step_avg:34.46ms
step:372/1600 train_time:12822ms step_avg:34.47ms
step:373/1600 train_time:12853ms step_avg:34.46ms
step:374/1600 train_time:12890ms step_avg:34.46ms
step:375/1600 train_time:12920ms step_avg:34.45ms
step:376/1600 train_time:12958ms step_avg:34.46ms
step:377/1600 train_time:12989ms step_avg:34.45ms
step:378/1600 train_time:13026ms step_avg:34.46ms
step:379/1600 train_time:13056ms step_avg:34.45ms
step:380/1600 train_time:13093ms step_avg:34.46ms
step:381/1600 train_time:13124ms step_avg:34.45ms
step:382/1600 train_time:13161ms step_avg:34.45ms
step:383/1600 train_time:13192ms step_avg:34.44ms
step:384/1600 train_time:13229ms step_avg:34.45ms
step:385/1600 train_time:13259ms step_avg:34.44ms
step:386/1600 train_time:13296ms step_avg:34.45ms
step:387/1600 train_time:13327ms step_avg:34.44ms
step:388/1600 train_time:13364ms step_avg:34.44ms
step:389/1600 train_time:13395ms step_avg:34.43ms
step:390/1600 train_time:13432ms step_avg:34.44ms
step:391/1600 train_time:13463ms step_avg:34.43ms
step:392/1600 train_time:13500ms step_avg:34.44ms
step:393/1600 train_time:13531ms step_avg:34.43ms
step:394/1600 train_time:13568ms step_avg:34.44ms
step:395/1600 train_time:13599ms step_avg:34.43ms
step:396/1600 train_time:13635ms step_avg:34.43ms
step:397/1600 train_time:13666ms step_avg:34.42ms
step:398/1600 train_time:13703ms step_avg:34.43ms
step:399/1600 train_time:13734ms step_avg:34.42ms
step:400/1600 train_time:13771ms step_avg:34.43ms
step:401/1600 train_time:13801ms step_avg:34.42ms
step:402/1600 train_time:13839ms step_avg:34.43ms
step:403/1600 train_time:13870ms step_avg:34.42ms
step:404/1600 train_time:13907ms step_avg:34.42ms
step:405/1600 train_time:13938ms step_avg:34.42ms
step:406/1600 train_time:13975ms step_avg:34.42ms
step:407/1600 train_time:14005ms step_avg:34.41ms
step:408/1600 train_time:14043ms step_avg:34.42ms
step:409/1600 train_time:14073ms step_avg:34.41ms
step:410/1600 train_time:14110ms step_avg:34.41ms
step:411/1600 train_time:14141ms step_avg:34.41ms
step:412/1600 train_time:14178ms step_avg:34.41ms
step:413/1600 train_time:14209ms step_avg:34.40ms
step:414/1600 train_time:14246ms step_avg:34.41ms
step:415/1600 train_time:14277ms step_avg:34.40ms
step:416/1600 train_time:14313ms step_avg:34.41ms
step:417/1600 train_time:14344ms step_avg:34.40ms
step:418/1600 train_time:14381ms step_avg:34.40ms
step:419/1600 train_time:14412ms step_avg:34.40ms
step:420/1600 train_time:14449ms step_avg:34.40ms
step:421/1600 train_time:14480ms step_avg:34.39ms
step:422/1600 train_time:14517ms step_avg:34.40ms
step:423/1600 train_time:14548ms step_avg:34.39ms
step:424/1600 train_time:14585ms step_avg:34.40ms
step:425/1600 train_time:14615ms step_avg:34.39ms
step:426/1600 train_time:14652ms step_avg:34.40ms
step:427/1600 train_time:14684ms step_avg:34.39ms
step:428/1600 train_time:14721ms step_avg:34.39ms
step:429/1600 train_time:14751ms step_avg:34.39ms
step:430/1600 train_time:14788ms step_avg:34.39ms
step:431/1600 train_time:14820ms step_avg:34.38ms
step:432/1600 train_time:14857ms step_avg:34.39ms
step:433/1600 train_time:14887ms step_avg:34.38ms
step:434/1600 train_time:14924ms step_avg:34.39ms
step:435/1600 train_time:14956ms step_avg:34.38ms
step:436/1600 train_time:14993ms step_avg:34.39ms
step:437/1600 train_time:15024ms step_avg:34.38ms
step:438/1600 train_time:15061ms step_avg:34.39ms
step:439/1600 train_time:15092ms step_avg:34.38ms
step:440/1600 train_time:15129ms step_avg:34.38ms
step:441/1600 train_time:15159ms step_avg:34.38ms
step:442/1600 train_time:15196ms step_avg:34.38ms
step:443/1600 train_time:15227ms step_avg:34.37ms
step:444/1600 train_time:15264ms step_avg:34.38ms
step:445/1600 train_time:15294ms step_avg:34.37ms
step:446/1600 train_time:15331ms step_avg:34.38ms
step:447/1600 train_time:15362ms step_avg:34.37ms
step:448/1600 train_time:15399ms step_avg:34.37ms
step:449/1600 train_time:15431ms step_avg:34.37ms
step:450/1600 train_time:15468ms step_avg:34.37ms
step:451/1600 train_time:15499ms step_avg:34.36ms
step:452/1600 train_time:15536ms step_avg:34.37ms
step:453/1600 train_time:15567ms step_avg:34.36ms
step:454/1600 train_time:15604ms step_avg:34.37ms
step:455/1600 train_time:15635ms step_avg:34.36ms
step:456/1600 train_time:15672ms step_avg:34.37ms
step:457/1600 train_time:15703ms step_avg:34.36ms
step:458/1600 train_time:15740ms step_avg:34.37ms
step:459/1600 train_time:15771ms step_avg:34.36ms
step:460/1600 train_time:15807ms step_avg:34.36ms
step:461/1600 train_time:15839ms step_avg:34.36ms
step:462/1600 train_time:15876ms step_avg:34.36ms
step:463/1600 train_time:15907ms step_avg:34.36ms
step:464/1600 train_time:15944ms step_avg:34.36ms
step:465/1600 train_time:15974ms step_avg:34.35ms
step:466/1600 train_time:16011ms step_avg:34.36ms
step:467/1600 train_time:16042ms step_avg:34.35ms
step:468/1600 train_time:16080ms step_avg:34.36ms
step:469/1600 train_time:16110ms step_avg:34.35ms
step:470/1600 train_time:16147ms step_avg:34.36ms
step:471/1600 train_time:16179ms step_avg:34.35ms
step:472/1600 train_time:16216ms step_avg:34.36ms
step:473/1600 train_time:16247ms step_avg:34.35ms
step:474/1600 train_time:16283ms step_avg:34.35ms
step:475/1600 train_time:16314ms step_avg:34.34ms
step:476/1600 train_time:16351ms step_avg:34.35ms
step:477/1600 train_time:16381ms step_avg:34.34ms
step:478/1600 train_time:16418ms step_avg:34.35ms
step:479/1600 train_time:16449ms step_avg:34.34ms
step:480/1600 train_time:16486ms step_avg:34.35ms
step:481/1600 train_time:16517ms step_avg:34.34ms
step:482/1600 train_time:16554ms step_avg:34.34ms
step:483/1600 train_time:16584ms step_avg:34.34ms
step:484/1600 train_time:16621ms step_avg:34.34ms
step:485/1600 train_time:16652ms step_avg:34.33ms
step:486/1600 train_time:16689ms step_avg:34.34ms
step:487/1600 train_time:16720ms step_avg:34.33ms
step:488/1600 train_time:16757ms step_avg:34.34ms
step:489/1600 train_time:16787ms step_avg:34.33ms
step:490/1600 train_time:16825ms step_avg:34.34ms
step:491/1600 train_time:16855ms step_avg:34.33ms
step:492/1600 train_time:16892ms step_avg:34.33ms
step:493/1600 train_time:16923ms step_avg:34.33ms
step:494/1600 train_time:16960ms step_avg:34.33ms
step:495/1600 train_time:16990ms step_avg:34.32ms
step:496/1600 train_time:17027ms step_avg:34.33ms
step:497/1600 train_time:17058ms step_avg:34.32ms
step:498/1600 train_time:17095ms step_avg:34.33ms
step:499/1600 train_time:17126ms step_avg:34.32ms
step:500/1600 train_time:17163ms step_avg:34.33ms
step:500/1600 val_loss:4.2440 train_time:17211ms step_avg:34.42ms
step:501/1600 train_time:17228ms step_avg:34.39ms
step:502/1600 train_time:17247ms step_avg:34.36ms
step:503/1600 train_time:17265ms step_avg:34.32ms
step:504/1600 train_time:17304ms step_avg:34.33ms
step:505/1600 train_time:17337ms step_avg:34.33ms
step:506/1600 train_time:17375ms step_avg:34.34ms
step:507/1600 train_time:17406ms step_avg:34.33ms
step:508/1600 train_time:17443ms step_avg:34.34ms
step:509/1600 train_time:17474ms step_avg:34.33ms
step:510/1600 train_time:17512ms step_avg:34.34ms
step:511/1600 train_time:17542ms step_avg:34.33ms
step:512/1600 train_time:17580ms step_avg:34.34ms
step:513/1600 train_time:17610ms step_avg:34.33ms
step:514/1600 train_time:17647ms step_avg:34.33ms
step:515/1600 train_time:17678ms step_avg:34.33ms
step:516/1600 train_time:17715ms step_avg:34.33ms
step:517/1600 train_time:17746ms step_avg:34.32ms
step:518/1600 train_time:17782ms step_avg:34.33ms
step:519/1600 train_time:17813ms step_avg:34.32ms
step:520/1600 train_time:17849ms step_avg:34.33ms
step:521/1600 train_time:17920ms step_avg:34.40ms
step:522/1600 train_time:17976ms step_avg:34.44ms
step:523/1600 train_time:18037ms step_avg:34.49ms
step:524/1600 train_time:18096ms step_avg:34.53ms
step:525/1600 train_time:18157ms step_avg:34.59ms
step:526/1600 train_time:18216ms step_avg:34.63ms
step:527/1600 train_time:18280ms step_avg:34.69ms
step:528/1600 train_time:18340ms step_avg:34.74ms
step:529/1600 train_time:18404ms step_avg:34.79ms
step:530/1600 train_time:18463ms step_avg:34.84ms
step:531/1600 train_time:18526ms step_avg:34.89ms
step:532/1600 train_time:18585ms step_avg:34.93ms
step:533/1600 train_time:18648ms step_avg:34.99ms
step:534/1600 train_time:18707ms step_avg:35.03ms
step:535/1600 train_time:18769ms step_avg:35.08ms
step:536/1600 train_time:18829ms step_avg:35.13ms
step:537/1600 train_time:18891ms step_avg:35.18ms
step:538/1600 train_time:18952ms step_avg:35.23ms
step:539/1600 train_time:19014ms step_avg:35.28ms
step:540/1600 train_time:19070ms step_avg:35.32ms
step:541/1600 train_time:19133ms step_avg:35.37ms
step:542/1600 train_time:19191ms step_avg:35.41ms
step:543/1600 train_time:19254ms step_avg:35.46ms
step:544/1600 train_time:19313ms step_avg:35.50ms
step:545/1600 train_time:19376ms step_avg:35.55ms
step:546/1600 train_time:19436ms step_avg:35.60ms
step:547/1600 train_time:19499ms step_avg:35.65ms
step:548/1600 train_time:19558ms step_avg:35.69ms
step:549/1600 train_time:19622ms step_avg:35.74ms
step:550/1600 train_time:19682ms step_avg:35.78ms
step:551/1600 train_time:19745ms step_avg:35.83ms
step:552/1600 train_time:19804ms step_avg:35.88ms
step:553/1600 train_time:19866ms step_avg:35.92ms
step:554/1600 train_time:19925ms step_avg:35.97ms
step:555/1600 train_time:19987ms step_avg:36.01ms
step:556/1600 train_time:20046ms step_avg:36.05ms
step:557/1600 train_time:20108ms step_avg:36.10ms
step:558/1600 train_time:20168ms step_avg:36.14ms
step:559/1600 train_time:20232ms step_avg:36.19ms
step:560/1600 train_time:20292ms step_avg:36.24ms
step:561/1600 train_time:20356ms step_avg:36.28ms
step:562/1600 train_time:20414ms step_avg:36.32ms
step:563/1600 train_time:20476ms step_avg:36.37ms
step:564/1600 train_time:20536ms step_avg:36.41ms
step:565/1600 train_time:20599ms step_avg:36.46ms
step:566/1600 train_time:20659ms step_avg:36.50ms
step:567/1600 train_time:20722ms step_avg:36.55ms
step:568/1600 train_time:20781ms step_avg:36.59ms
step:569/1600 train_time:20844ms step_avg:36.63ms
step:570/1600 train_time:20903ms step_avg:36.67ms
step:571/1600 train_time:20966ms step_avg:36.72ms
step:572/1600 train_time:21025ms step_avg:36.76ms
step:573/1600 train_time:21092ms step_avg:36.81ms
step:574/1600 train_time:21149ms step_avg:36.85ms
step:575/1600 train_time:21211ms step_avg:36.89ms
step:576/1600 train_time:21269ms step_avg:36.93ms
step:577/1600 train_time:21333ms step_avg:36.97ms
step:578/1600 train_time:21391ms step_avg:37.01ms
step:579/1600 train_time:21456ms step_avg:37.06ms
step:580/1600 train_time:21514ms step_avg:37.09ms
step:581/1600 train_time:21576ms step_avg:37.14ms
step:582/1600 train_time:21635ms step_avg:37.17ms
step:583/1600 train_time:21698ms step_avg:37.22ms
step:584/1600 train_time:21757ms step_avg:37.25ms
step:585/1600 train_time:21820ms step_avg:37.30ms
step:586/1600 train_time:21879ms step_avg:37.34ms
step:587/1600 train_time:21942ms step_avg:37.38ms
step:588/1600 train_time:22001ms step_avg:37.42ms
step:589/1600 train_time:22064ms step_avg:37.46ms
step:590/1600 train_time:22124ms step_avg:37.50ms
step:591/1600 train_time:22186ms step_avg:37.54ms
step:592/1600 train_time:22245ms step_avg:37.58ms
step:593/1600 train_time:22308ms step_avg:37.62ms
step:594/1600 train_time:22366ms step_avg:37.65ms
step:595/1600 train_time:22428ms step_avg:37.69ms
step:596/1600 train_time:22487ms step_avg:37.73ms
step:597/1600 train_time:22551ms step_avg:37.77ms
step:598/1600 train_time:22612ms step_avg:37.81ms
step:599/1600 train_time:22673ms step_avg:37.85ms
step:600/1600 train_time:22733ms step_avg:37.89ms
step:601/1600 train_time:22795ms step_avg:37.93ms
step:602/1600 train_time:22854ms step_avg:37.96ms
step:603/1600 train_time:22916ms step_avg:38.00ms
step:604/1600 train_time:22975ms step_avg:38.04ms
step:605/1600 train_time:23040ms step_avg:38.08ms
step:606/1600 train_time:23098ms step_avg:38.12ms
step:607/1600 train_time:23161ms step_avg:38.16ms
step:608/1600 train_time:23220ms step_avg:38.19ms
step:609/1600 train_time:23283ms step_avg:38.23ms
step:610/1600 train_time:23342ms step_avg:38.27ms
step:611/1600 train_time:23405ms step_avg:38.31ms
step:612/1600 train_time:23464ms step_avg:38.34ms
step:613/1600 train_time:23527ms step_avg:38.38ms
step:614/1600 train_time:23585ms step_avg:38.41ms
step:615/1600 train_time:23647ms step_avg:38.45ms
step:616/1600 train_time:23706ms step_avg:38.48ms
step:617/1600 train_time:23769ms step_avg:38.52ms
step:618/1600 train_time:23828ms step_avg:38.56ms
step:619/1600 train_time:23890ms step_avg:38.60ms
step:620/1600 train_time:23951ms step_avg:38.63ms
step:621/1600 train_time:24013ms step_avg:38.67ms
step:622/1600 train_time:24074ms step_avg:38.70ms
step:623/1600 train_time:24137ms step_avg:38.74ms
step:624/1600 train_time:24195ms step_avg:38.77ms
step:625/1600 train_time:24258ms step_avg:38.81ms
step:626/1600 train_time:24318ms step_avg:38.85ms
step:627/1600 train_time:24381ms step_avg:38.88ms
step:628/1600 train_time:24440ms step_avg:38.92ms
step:629/1600 train_time:24503ms step_avg:38.96ms
step:630/1600 train_time:24562ms step_avg:38.99ms
step:631/1600 train_time:24625ms step_avg:39.02ms
step:632/1600 train_time:24683ms step_avg:39.06ms
step:633/1600 train_time:24746ms step_avg:39.09ms
step:634/1600 train_time:24804ms step_avg:39.12ms
step:635/1600 train_time:24866ms step_avg:39.16ms
step:636/1600 train_time:24924ms step_avg:39.19ms
step:637/1600 train_time:24987ms step_avg:39.23ms
step:638/1600 train_time:25047ms step_avg:39.26ms
step:639/1600 train_time:25110ms step_avg:39.30ms
step:640/1600 train_time:25170ms step_avg:39.33ms
step:641/1600 train_time:25233ms step_avg:39.37ms
step:642/1600 train_time:25292ms step_avg:39.40ms
step:643/1600 train_time:25355ms step_avg:39.43ms
step:644/1600 train_time:25414ms step_avg:39.46ms
step:645/1600 train_time:25477ms step_avg:39.50ms
step:646/1600 train_time:25536ms step_avg:39.53ms
step:647/1600 train_time:25599ms step_avg:39.56ms
step:648/1600 train_time:25658ms step_avg:39.59ms
step:649/1600 train_time:25721ms step_avg:39.63ms
step:650/1600 train_time:25780ms step_avg:39.66ms
step:651/1600 train_time:25843ms step_avg:39.70ms
step:652/1600 train_time:25903ms step_avg:39.73ms
step:653/1600 train_time:25966ms step_avg:39.76ms
step:654/1600 train_time:26024ms step_avg:39.79ms
step:655/1600 train_time:26086ms step_avg:39.83ms
step:656/1600 train_time:26147ms step_avg:39.86ms
step:657/1600 train_time:26209ms step_avg:39.89ms
step:658/1600 train_time:26267ms step_avg:39.92ms
step:659/1600 train_time:26330ms step_avg:39.95ms
step:660/1600 train_time:26389ms step_avg:39.98ms
step:661/1600 train_time:26453ms step_avg:40.02ms
step:662/1600 train_time:26512ms step_avg:40.05ms
step:663/1600 train_time:26575ms step_avg:40.08ms
step:664/1600 train_time:26633ms step_avg:40.11ms
step:665/1600 train_time:26696ms step_avg:40.14ms
step:666/1600 train_time:26755ms step_avg:40.17ms
step:667/1600 train_time:26817ms step_avg:40.21ms
step:668/1600 train_time:26877ms step_avg:40.23ms
step:669/1600 train_time:26941ms step_avg:40.27ms
step:670/1600 train_time:27000ms step_avg:40.30ms
step:671/1600 train_time:27063ms step_avg:40.33ms
step:672/1600 train_time:27122ms step_avg:40.36ms
step:673/1600 train_time:27188ms step_avg:40.40ms
step:674/1600 train_time:27244ms step_avg:40.42ms
step:675/1600 train_time:27306ms step_avg:40.45ms
step:676/1600 train_time:27365ms step_avg:40.48ms
step:677/1600 train_time:27428ms step_avg:40.51ms
step:678/1600 train_time:27486ms step_avg:40.54ms
step:679/1600 train_time:27549ms step_avg:40.57ms
step:680/1600 train_time:27608ms step_avg:40.60ms
step:681/1600 train_time:27673ms step_avg:40.64ms
step:682/1600 train_time:27732ms step_avg:40.66ms
step:683/1600 train_time:27794ms step_avg:40.69ms
step:684/1600 train_time:27853ms step_avg:40.72ms
step:685/1600 train_time:27915ms step_avg:40.75ms
step:686/1600 train_time:27975ms step_avg:40.78ms
step:687/1600 train_time:28037ms step_avg:40.81ms
step:688/1600 train_time:28097ms step_avg:40.84ms
step:689/1600 train_time:28159ms step_avg:40.87ms
step:690/1600 train_time:28218ms step_avg:40.90ms
step:691/1600 train_time:28282ms step_avg:40.93ms
step:692/1600 train_time:28342ms step_avg:40.96ms
step:693/1600 train_time:28405ms step_avg:40.99ms
step:694/1600 train_time:28463ms step_avg:41.01ms
step:695/1600 train_time:28527ms step_avg:41.05ms
step:696/1600 train_time:28585ms step_avg:41.07ms
step:697/1600 train_time:28648ms step_avg:41.10ms
step:698/1600 train_time:28707ms step_avg:41.13ms
step:699/1600 train_time:28769ms step_avg:41.16ms
step:700/1600 train_time:28829ms step_avg:41.18ms
step:701/1600 train_time:28891ms step_avg:41.21ms
step:702/1600 train_time:28951ms step_avg:41.24ms
step:703/1600 train_time:29014ms step_avg:41.27ms
step:704/1600 train_time:29073ms step_avg:41.30ms
step:705/1600 train_time:29135ms step_avg:41.33ms
step:706/1600 train_time:29196ms step_avg:41.35ms
step:707/1600 train_time:29258ms step_avg:41.38ms
step:708/1600 train_time:29316ms step_avg:41.41ms
step:709/1600 train_time:29380ms step_avg:41.44ms
step:710/1600 train_time:29439ms step_avg:41.46ms
step:711/1600 train_time:29502ms step_avg:41.49ms
step:712/1600 train_time:29562ms step_avg:41.52ms
step:713/1600 train_time:29624ms step_avg:41.55ms
step:714/1600 train_time:29683ms step_avg:41.57ms
step:715/1600 train_time:29746ms step_avg:41.60ms
step:716/1600 train_time:29804ms step_avg:41.63ms
step:717/1600 train_time:29867ms step_avg:41.66ms
step:718/1600 train_time:29927ms step_avg:41.68ms
step:719/1600 train_time:29989ms step_avg:41.71ms
step:720/1600 train_time:30048ms step_avg:41.73ms
step:721/1600 train_time:30111ms step_avg:41.76ms
step:722/1600 train_time:30171ms step_avg:41.79ms
step:723/1600 train_time:30233ms step_avg:41.82ms
step:724/1600 train_time:30293ms step_avg:41.84ms
step:725/1600 train_time:30356ms step_avg:41.87ms
step:726/1600 train_time:30416ms step_avg:41.89ms
step:727/1600 train_time:30478ms step_avg:41.92ms
step:728/1600 train_time:30537ms step_avg:41.95ms
step:729/1600 train_time:30600ms step_avg:41.98ms
step:730/1600 train_time:30660ms step_avg:42.00ms
step:731/1600 train_time:30723ms step_avg:42.03ms
step:732/1600 train_time:30782ms step_avg:42.05ms
step:733/1600 train_time:30845ms step_avg:42.08ms
step:734/1600 train_time:30904ms step_avg:42.10ms
step:735/1600 train_time:30969ms step_avg:42.14ms
step:736/1600 train_time:31028ms step_avg:42.16ms
step:737/1600 train_time:31088ms step_avg:42.18ms
step:738/1600 train_time:31147ms step_avg:42.21ms
step:739/1600 train_time:31210ms step_avg:42.23ms
step:740/1600 train_time:31269ms step_avg:42.26ms
step:741/1600 train_time:31333ms step_avg:42.28ms
step:742/1600 train_time:31392ms step_avg:42.31ms
step:743/1600 train_time:31455ms step_avg:42.34ms
step:744/1600 train_time:31514ms step_avg:42.36ms
step:745/1600 train_time:31576ms step_avg:42.38ms
step:746/1600 train_time:31636ms step_avg:42.41ms
step:747/1600 train_time:31700ms step_avg:42.44ms
step:748/1600 train_time:31758ms step_avg:42.46ms
step:749/1600 train_time:31821ms step_avg:42.48ms
step:750/1600 train_time:31880ms step_avg:42.51ms
step:750/1600 val_loss:3.8896 train_time:31927ms step_avg:42.57ms
step:751/1600 train_time:31946ms step_avg:42.54ms
step:752/1600 train_time:32003ms step_avg:42.56ms
step:753/1600 train_time:32068ms step_avg:42.59ms
step:754/1600 train_time:32129ms step_avg:42.61ms
step:755/1600 train_time:32191ms step_avg:42.64ms
step:756/1600 train_time:32249ms step_avg:42.66ms
step:757/1600 train_time:32311ms step_avg:42.68ms
step:758/1600 train_time:32371ms step_avg:42.71ms
step:759/1600 train_time:32432ms step_avg:42.73ms
step:760/1600 train_time:32491ms step_avg:42.75ms
step:761/1600 train_time:32553ms step_avg:42.78ms
step:762/1600 train_time:32611ms step_avg:42.80ms
step:763/1600 train_time:32672ms step_avg:42.82ms
step:764/1600 train_time:32731ms step_avg:42.84ms
step:765/1600 train_time:32792ms step_avg:42.86ms
step:766/1600 train_time:32850ms step_avg:42.89ms
step:767/1600 train_time:32913ms step_avg:42.91ms
step:768/1600 train_time:32973ms step_avg:42.93ms
step:769/1600 train_time:33038ms step_avg:42.96ms
step:770/1600 train_time:33098ms step_avg:42.99ms
step:771/1600 train_time:33163ms step_avg:43.01ms
step:772/1600 train_time:33221ms step_avg:43.03ms
step:773/1600 train_time:33283ms step_avg:43.06ms
step:774/1600 train_time:33342ms step_avg:43.08ms
step:775/1600 train_time:33404ms step_avg:43.10ms
step:776/1600 train_time:33463ms step_avg:43.12ms
step:777/1600 train_time:33525ms step_avg:43.15ms
step:778/1600 train_time:33584ms step_avg:43.17ms
step:779/1600 train_time:33646ms step_avg:43.19ms
step:780/1600 train_time:33704ms step_avg:43.21ms
step:781/1600 train_time:33767ms step_avg:43.24ms
step:782/1600 train_time:33825ms step_avg:43.25ms
step:783/1600 train_time:33888ms step_avg:43.28ms
step:784/1600 train_time:33948ms step_avg:43.30ms
step:785/1600 train_time:34012ms step_avg:43.33ms
step:786/1600 train_time:34071ms step_avg:43.35ms
step:787/1600 train_time:34134ms step_avg:43.37ms
step:788/1600 train_time:34193ms step_avg:43.39ms
step:789/1600 train_time:34255ms step_avg:43.42ms
step:790/1600 train_time:34314ms step_avg:43.44ms
step:791/1600 train_time:34377ms step_avg:43.46ms
step:792/1600 train_time:34435ms step_avg:43.48ms
step:793/1600 train_time:34497ms step_avg:43.50ms
step:794/1600 train_time:34557ms step_avg:43.52ms
step:795/1600 train_time:34619ms step_avg:43.55ms
step:796/1600 train_time:34678ms step_avg:43.57ms
step:797/1600 train_time:34740ms step_avg:43.59ms
step:798/1600 train_time:34800ms step_avg:43.61ms
step:799/1600 train_time:34862ms step_avg:43.63ms
step:800/1600 train_time:34921ms step_avg:43.65ms
step:801/1600 train_time:34985ms step_avg:43.68ms
step:802/1600 train_time:35044ms step_avg:43.70ms
step:803/1600 train_time:35107ms step_avg:43.72ms
step:804/1600 train_time:35166ms step_avg:43.74ms
step:805/1600 train_time:35228ms step_avg:43.76ms
step:806/1600 train_time:35288ms step_avg:43.78ms
step:807/1600 train_time:35350ms step_avg:43.80ms
step:808/1600 train_time:35409ms step_avg:43.82ms
step:809/1600 train_time:35472ms step_avg:43.85ms
step:810/1600 train_time:35531ms step_avg:43.87ms
step:811/1600 train_time:35593ms step_avg:43.89ms
step:812/1600 train_time:35651ms step_avg:43.91ms
step:813/1600 train_time:35714ms step_avg:43.93ms
step:814/1600 train_time:35772ms step_avg:43.95ms
step:815/1600 train_time:35835ms step_avg:43.97ms
step:816/1600 train_time:35894ms step_avg:43.99ms
step:817/1600 train_time:35957ms step_avg:44.01ms
step:818/1600 train_time:36017ms step_avg:44.03ms
step:819/1600 train_time:36080ms step_avg:44.05ms
step:820/1600 train_time:36139ms step_avg:44.07ms
step:821/1600 train_time:36202ms step_avg:44.09ms
step:822/1600 train_time:36261ms step_avg:44.11ms
step:823/1600 train_time:36322ms step_avg:44.13ms
step:824/1600 train_time:36382ms step_avg:44.15ms
step:825/1600 train_time:36444ms step_avg:44.18ms
step:826/1600 train_time:36503ms step_avg:44.19ms
step:827/1600 train_time:36565ms step_avg:44.21ms
step:828/1600 train_time:36625ms step_avg:44.23ms
step:829/1600 train_time:36687ms step_avg:44.26ms
step:830/1600 train_time:36747ms step_avg:44.27ms
step:831/1600 train_time:36809ms step_avg:44.30ms
step:832/1600 train_time:36869ms step_avg:44.31ms
step:833/1600 train_time:36932ms step_avg:44.34ms
step:834/1600 train_time:36991ms step_avg:44.35ms
step:835/1600 train_time:37053ms step_avg:44.38ms
step:836/1600 train_time:37112ms step_avg:44.39ms
step:837/1600 train_time:37176ms step_avg:44.42ms
step:838/1600 train_time:37234ms step_avg:44.43ms
step:839/1600 train_time:37296ms step_avg:44.45ms
step:840/1600 train_time:37355ms step_avg:44.47ms
step:841/1600 train_time:37418ms step_avg:44.49ms
step:842/1600 train_time:37477ms step_avg:44.51ms
step:843/1600 train_time:37540ms step_avg:44.53ms
step:844/1600 train_time:37599ms step_avg:44.55ms
step:845/1600 train_time:37661ms step_avg:44.57ms
step:846/1600 train_time:37720ms step_avg:44.59ms
step:847/1600 train_time:37782ms step_avg:44.61ms
step:848/1600 train_time:37841ms step_avg:44.62ms
step:849/1600 train_time:37904ms step_avg:44.65ms
step:850/1600 train_time:37963ms step_avg:44.66ms
step:851/1600 train_time:38026ms step_avg:44.68ms
step:852/1600 train_time:38085ms step_avg:44.70ms
step:853/1600 train_time:38148ms step_avg:44.72ms
step:854/1600 train_time:38207ms step_avg:44.74ms
step:855/1600 train_time:38269ms step_avg:44.76ms
step:856/1600 train_time:38328ms step_avg:44.78ms
step:857/1600 train_time:38391ms step_avg:44.80ms
step:858/1600 train_time:38450ms step_avg:44.81ms
step:859/1600 train_time:38512ms step_avg:44.83ms
step:860/1600 train_time:38571ms step_avg:44.85ms
step:861/1600 train_time:38633ms step_avg:44.87ms
step:862/1600 train_time:38692ms step_avg:44.89ms
step:863/1600 train_time:38755ms step_avg:44.91ms
step:864/1600 train_time:38814ms step_avg:44.92ms
step:865/1600 train_time:38877ms step_avg:44.94ms
step:866/1600 train_time:38936ms step_avg:44.96ms
step:867/1600 train_time:38998ms step_avg:44.98ms
step:868/1600 train_time:39058ms step_avg:45.00ms
step:869/1600 train_time:39121ms step_avg:45.02ms
step:870/1600 train_time:39180ms step_avg:45.03ms
step:871/1600 train_time:39243ms step_avg:45.05ms
step:872/1600 train_time:39302ms step_avg:45.07ms
step:873/1600 train_time:39365ms step_avg:45.09ms
step:874/1600 train_time:39424ms step_avg:45.11ms
step:875/1600 train_time:39488ms step_avg:45.13ms
step:876/1600 train_time:39546ms step_avg:45.14ms
step:877/1600 train_time:39608ms step_avg:45.16ms
step:878/1600 train_time:39669ms step_avg:45.18ms
step:879/1600 train_time:39731ms step_avg:45.20ms
step:880/1600 train_time:39790ms step_avg:45.22ms
step:881/1600 train_time:39852ms step_avg:45.24ms
step:882/1600 train_time:39911ms step_avg:45.25ms
step:883/1600 train_time:39974ms step_avg:45.27ms
step:884/1600 train_time:40034ms step_avg:45.29ms
step:885/1600 train_time:40096ms step_avg:45.31ms
step:886/1600 train_time:40154ms step_avg:45.32ms
step:887/1600 train_time:40219ms step_avg:45.34ms
step:888/1600 train_time:40278ms step_avg:45.36ms
step:889/1600 train_time:40340ms step_avg:45.38ms
step:890/1600 train_time:40403ms step_avg:45.40ms
step:891/1600 train_time:40464ms step_avg:45.41ms
step:892/1600 train_time:40525ms step_avg:45.43ms
step:893/1600 train_time:40586ms step_avg:45.45ms
step:894/1600 train_time:40645ms step_avg:45.46ms
step:895/1600 train_time:40708ms step_avg:45.48ms
step:896/1600 train_time:40767ms step_avg:45.50ms
step:897/1600 train_time:40829ms step_avg:45.52ms
step:898/1600 train_time:40888ms step_avg:45.53ms
step:899/1600 train_time:40951ms step_avg:45.55ms
step:900/1600 train_time:41010ms step_avg:45.57ms
step:901/1600 train_time:41073ms step_avg:45.59ms
step:902/1600 train_time:41132ms step_avg:45.60ms
step:903/1600 train_time:41195ms step_avg:45.62ms
step:904/1600 train_time:41254ms step_avg:45.63ms
step:905/1600 train_time:41316ms step_avg:45.65ms
step:906/1600 train_time:41375ms step_avg:45.67ms
step:907/1600 train_time:41437ms step_avg:45.69ms
step:908/1600 train_time:41497ms step_avg:45.70ms
step:909/1600 train_time:41560ms step_avg:45.72ms
step:910/1600 train_time:41620ms step_avg:45.74ms
step:911/1600 train_time:41684ms step_avg:45.76ms
step:912/1600 train_time:41742ms step_avg:45.77ms
step:913/1600 train_time:41805ms step_avg:45.79ms
step:914/1600 train_time:41864ms step_avg:45.80ms
step:915/1600 train_time:41926ms step_avg:45.82ms
step:916/1600 train_time:41985ms step_avg:45.84ms
step:917/1600 train_time:42048ms step_avg:45.85ms
step:918/1600 train_time:42107ms step_avg:45.87ms
step:919/1600 train_time:42170ms step_avg:45.89ms
step:920/1600 train_time:42229ms step_avg:45.90ms
step:921/1600 train_time:42292ms step_avg:45.92ms
step:922/1600 train_time:42351ms step_avg:45.93ms
step:923/1600 train_time:42414ms step_avg:45.95ms
step:924/1600 train_time:42474ms step_avg:45.97ms
step:925/1600 train_time:42536ms step_avg:45.99ms
step:926/1600 train_time:42595ms step_avg:46.00ms
step:927/1600 train_time:42657ms step_avg:46.02ms
step:928/1600 train_time:42717ms step_avg:46.03ms
step:929/1600 train_time:42780ms step_avg:46.05ms
step:930/1600 train_time:42839ms step_avg:46.06ms
step:931/1600 train_time:42902ms step_avg:46.08ms
step:932/1600 train_time:42961ms step_avg:46.10ms
step:933/1600 train_time:43023ms step_avg:46.11ms
step:934/1600 train_time:43082ms step_avg:46.13ms
step:935/1600 train_time:43145ms step_avg:46.14ms
step:936/1600 train_time:43204ms step_avg:46.16ms
step:937/1600 train_time:43267ms step_avg:46.18ms
step:938/1600 train_time:43326ms step_avg:46.19ms
step:939/1600 train_time:43388ms step_avg:46.21ms
step:940/1600 train_time:43447ms step_avg:46.22ms
step:941/1600 train_time:43510ms step_avg:46.24ms
step:942/1600 train_time:43570ms step_avg:46.25ms
step:943/1600 train_time:43633ms step_avg:46.27ms
step:944/1600 train_time:43692ms step_avg:46.28ms
step:945/1600 train_time:43754ms step_avg:46.30ms
step:946/1600 train_time:43813ms step_avg:46.31ms
step:947/1600 train_time:43876ms step_avg:46.33ms
step:948/1600 train_time:43934ms step_avg:46.34ms
step:949/1600 train_time:43997ms step_avg:46.36ms
step:950/1600 train_time:44056ms step_avg:46.38ms
step:951/1600 train_time:44120ms step_avg:46.39ms
step:952/1600 train_time:44180ms step_avg:46.41ms
step:953/1600 train_time:44242ms step_avg:46.42ms
step:954/1600 train_time:44302ms step_avg:46.44ms
step:955/1600 train_time:44364ms step_avg:46.45ms
step:956/1600 train_time:44423ms step_avg:46.47ms
step:957/1600 train_time:44485ms step_avg:46.48ms
step:958/1600 train_time:44544ms step_avg:46.50ms
step:959/1600 train_time:44608ms step_avg:46.51ms
step:960/1600 train_time:44669ms step_avg:46.53ms
step:961/1600 train_time:44730ms step_avg:46.55ms
step:962/1600 train_time:44789ms step_avg:46.56ms
step:963/1600 train_time:44851ms step_avg:46.57ms
step:964/1600 train_time:44911ms step_avg:46.59ms
step:965/1600 train_time:44973ms step_avg:46.60ms
step:966/1600 train_time:45032ms step_avg:46.62ms
step:967/1600 train_time:45096ms step_avg:46.64ms
step:968/1600 train_time:45156ms step_avg:46.65ms
step:969/1600 train_time:45217ms step_avg:46.66ms
step:970/1600 train_time:45276ms step_avg:46.68ms
step:971/1600 train_time:45338ms step_avg:46.69ms
step:972/1600 train_time:45399ms step_avg:46.71ms
step:973/1600 train_time:45461ms step_avg:46.72ms
step:974/1600 train_time:45520ms step_avg:46.73ms
step:975/1600 train_time:45583ms step_avg:46.75ms
step:976/1600 train_time:45643ms step_avg:46.77ms
step:977/1600 train_time:45706ms step_avg:46.78ms
step:978/1600 train_time:45764ms step_avg:46.79ms
step:979/1600 train_time:45827ms step_avg:46.81ms
step:980/1600 train_time:45886ms step_avg:46.82ms
step:981/1600 train_time:45949ms step_avg:46.84ms
step:982/1600 train_time:46009ms step_avg:46.85ms
step:983/1600 train_time:46071ms step_avg:46.87ms
step:984/1600 train_time:46131ms step_avg:46.88ms
step:985/1600 train_time:46193ms step_avg:46.90ms
step:986/1600 train_time:46251ms step_avg:46.91ms
step:987/1600 train_time:46313ms step_avg:46.92ms
step:988/1600 train_time:46373ms step_avg:46.94ms
step:989/1600 train_time:46435ms step_avg:46.95ms
step:990/1600 train_time:46494ms step_avg:46.96ms
step:991/1600 train_time:46557ms step_avg:46.98ms
step:992/1600 train_time:46617ms step_avg:46.99ms
step:993/1600 train_time:46680ms step_avg:47.01ms
step:994/1600 train_time:46739ms step_avg:47.02ms
step:995/1600 train_time:46802ms step_avg:47.04ms
step:996/1600 train_time:46860ms step_avg:47.05ms
step:997/1600 train_time:46922ms step_avg:47.06ms
step:998/1600 train_time:46981ms step_avg:47.08ms
step:999/1600 train_time:47045ms step_avg:47.09ms
step:1000/1600 train_time:47105ms step_avg:47.10ms
step:1000/1600 val_loss:3.5989 train_time:47152ms step_avg:47.15ms
step:1001/1600 train_time:47172ms step_avg:47.13ms
step:1002/1600 train_time:47228ms step_avg:47.13ms
step:1003/1600 train_time:47293ms step_avg:47.15ms
step:1004/1600 train_time:47353ms step_avg:47.16ms
step:1005/1600 train_time:47415ms step_avg:47.18ms
step:1006/1600 train_time:47475ms step_avg:47.19ms
step:1007/1600 train_time:47537ms step_avg:47.21ms
step:1008/1600 train_time:47595ms step_avg:47.22ms
step:1009/1600 train_time:47657ms step_avg:47.23ms
step:1010/1600 train_time:47715ms step_avg:47.24ms
step:1011/1600 train_time:47777ms step_avg:47.26ms
step:1012/1600 train_time:47835ms step_avg:47.27ms
step:1013/1600 train_time:47896ms step_avg:47.28ms
step:1014/1600 train_time:47954ms step_avg:47.29ms
step:1015/1600 train_time:48016ms step_avg:47.31ms
step:1016/1600 train_time:48075ms step_avg:47.32ms
step:1017/1600 train_time:48139ms step_avg:47.33ms
step:1018/1600 train_time:48201ms step_avg:47.35ms
step:1019/1600 train_time:48266ms step_avg:47.37ms
step:1020/1600 train_time:48325ms step_avg:47.38ms
step:1021/1600 train_time:48386ms step_avg:47.39ms
step:1022/1600 train_time:48445ms step_avg:47.40ms
step:1023/1600 train_time:48507ms step_avg:47.42ms
step:1024/1600 train_time:48567ms step_avg:47.43ms
step:1025/1600 train_time:48630ms step_avg:47.44ms
step:1026/1600 train_time:48689ms step_avg:47.45ms
step:1027/1600 train_time:48751ms step_avg:47.47ms
step:1028/1600 train_time:48809ms step_avg:47.48ms
step:1029/1600 train_time:48871ms step_avg:47.49ms
step:1030/1600 train_time:48930ms step_avg:47.50ms
step:1031/1600 train_time:48992ms step_avg:47.52ms
step:1032/1600 train_time:49051ms step_avg:47.53ms
step:1033/1600 train_time:49114ms step_avg:47.54ms
step:1034/1600 train_time:49173ms step_avg:47.56ms
step:1035/1600 train_time:49236ms step_avg:47.57ms
step:1036/1600 train_time:49295ms step_avg:47.58ms
step:1037/1600 train_time:49358ms step_avg:47.60ms
step:1038/1600 train_time:49418ms step_avg:47.61ms
step:1039/1600 train_time:49480ms step_avg:47.62ms
step:1040/1600 train_time:49538ms step_avg:47.63ms
step:1041/1600 train_time:49609ms step_avg:47.66ms
step:1042/1600 train_time:49693ms step_avg:47.69ms
step:1043/1600 train_time:49781ms step_avg:47.73ms
step:1044/1600 train_time:49867ms step_avg:47.77ms
step:1045/1600 train_time:49955ms step_avg:47.80ms
step:1046/1600 train_time:50040ms step_avg:47.84ms
step:1047/1600 train_time:50128ms step_avg:47.88ms
step:1048/1600 train_time:50213ms step_avg:47.91ms
step:1049/1600 train_time:50303ms step_avg:47.95ms
step:1050/1600 train_time:50388ms step_avg:47.99ms
step:1051/1600 train_time:50476ms step_avg:48.03ms
step:1052/1600 train_time:50562ms step_avg:48.06ms
step:1053/1600 train_time:50650ms step_avg:48.10ms
step:1054/1600 train_time:50735ms step_avg:48.14ms
step:1055/1600 train_time:50824ms step_avg:48.17ms
step:1056/1600 train_time:50908ms step_avg:48.21ms
step:1057/1600 train_time:50997ms step_avg:48.25ms
step:1058/1600 train_time:51082ms step_avg:48.28ms
step:1059/1600 train_time:51171ms step_avg:48.32ms
step:1060/1600 train_time:51256ms step_avg:48.35ms
step:1061/1600 train_time:51345ms step_avg:48.39ms
step:1062/1600 train_time:51429ms step_avg:48.43ms
step:1063/1600 train_time:51517ms step_avg:48.46ms
step:1064/1600 train_time:51603ms step_avg:48.50ms
step:1065/1600 train_time:51691ms step_avg:48.54ms
step:1066/1600 train_time:51775ms step_avg:48.57ms
step:1067/1600 train_time:51863ms step_avg:48.61ms
step:1068/1600 train_time:51948ms step_avg:48.64ms
step:1069/1600 train_time:52036ms step_avg:48.68ms
step:1070/1600 train_time:52120ms step_avg:48.71ms
step:1071/1600 train_time:52208ms step_avg:48.75ms
step:1072/1600 train_time:52294ms step_avg:48.78ms
step:1073/1600 train_time:52382ms step_avg:48.82ms
step:1074/1600 train_time:52468ms step_avg:48.85ms
step:1075/1600 train_time:52557ms step_avg:48.89ms
step:1076/1600 train_time:52641ms step_avg:48.92ms
step:1077/1600 train_time:52731ms step_avg:48.96ms
step:1078/1600 train_time:52815ms step_avg:48.99ms
step:1079/1600 train_time:52903ms step_avg:49.03ms
step:1080/1600 train_time:52988ms step_avg:49.06ms
step:1081/1600 train_time:53076ms step_avg:49.10ms
step:1082/1600 train_time:53161ms step_avg:49.13ms
step:1083/1600 train_time:53250ms step_avg:49.17ms
step:1084/1600 train_time:53334ms step_avg:49.20ms
step:1085/1600 train_time:53423ms step_avg:49.24ms
step:1086/1600 train_time:53509ms step_avg:49.27ms
step:1087/1600 train_time:53597ms step_avg:49.31ms
step:1088/1600 train_time:53683ms step_avg:49.34ms
step:1089/1600 train_time:53771ms step_avg:49.38ms
step:1090/1600 train_time:53856ms step_avg:49.41ms
step:1091/1600 train_time:53944ms step_avg:49.44ms
step:1092/1600 train_time:54029ms step_avg:49.48ms
step:1093/1600 train_time:54117ms step_avg:49.51ms
step:1094/1600 train_time:54202ms step_avg:49.54ms
step:1095/1600 train_time:54290ms step_avg:49.58ms
step:1096/1600 train_time:54377ms step_avg:49.61ms
step:1097/1600 train_time:54465ms step_avg:49.65ms
step:1098/1600 train_time:54550ms step_avg:49.68ms
step:1099/1600 train_time:54639ms step_avg:49.72ms
step:1100/1600 train_time:54725ms step_avg:49.75ms
step:1101/1600 train_time:54813ms step_avg:49.78ms
step:1102/1600 train_time:54898ms step_avg:49.82ms
step:1103/1600 train_time:54986ms step_avg:49.85ms
step:1104/1600 train_time:55071ms step_avg:49.88ms
step:1105/1600 train_time:55160ms step_avg:49.92ms
step:1106/1600 train_time:55245ms step_avg:49.95ms
step:1107/1600 train_time:55334ms step_avg:49.99ms
step:1108/1600 train_time:55419ms step_avg:50.02ms
step:1109/1600 train_time:55507ms step_avg:50.05ms
step:1110/1600 train_time:55593ms step_avg:50.08ms
step:1111/1600 train_time:55681ms step_avg:50.12ms
step:1112/1600 train_time:55766ms step_avg:50.15ms
step:1113/1600 train_time:55855ms step_avg:50.18ms
step:1114/1600 train_time:55939ms step_avg:50.21ms
step:1115/1600 train_time:56027ms step_avg:50.25ms
step:1116/1600 train_time:56112ms step_avg:50.28ms
step:1117/1600 train_time:56200ms step_avg:50.31ms
step:1118/1600 train_time:56285ms step_avg:50.34ms
step:1119/1600 train_time:56373ms step_avg:50.38ms
step:1120/1600 train_time:56458ms step_avg:50.41ms
step:1121/1600 train_time:56547ms step_avg:50.44ms
step:1122/1600 train_time:56633ms step_avg:50.48ms
step:1123/1600 train_time:56722ms step_avg:50.51ms
step:1124/1600 train_time:56811ms step_avg:50.54ms
step:1125/1600 train_time:56895ms step_avg:50.57ms
step:1126/1600 train_time:56981ms step_avg:50.60ms
step:1127/1600 train_time:57071ms step_avg:50.64ms
step:1128/1600 train_time:57158ms step_avg:50.67ms
step:1129/1600 train_time:57248ms step_avg:50.71ms
step:1130/1600 train_time:57330ms step_avg:50.73ms
step:1131/1600 train_time:57414ms step_avg:50.76ms
step:1132/1600 train_time:57499ms step_avg:50.79ms
step:1133/1600 train_time:57589ms step_avg:50.83ms
step:1134/1600 train_time:57674ms step_avg:50.86ms
step:1135/1600 train_time:57762ms step_avg:50.89ms
step:1136/1600 train_time:57847ms step_avg:50.92ms
step:1137/1600 train_time:57936ms step_avg:50.96ms
step:1138/1600 train_time:58021ms step_avg:50.99ms
step:1139/1600 train_time:58109ms step_avg:51.02ms
step:1140/1600 train_time:58195ms step_avg:51.05ms
step:1141/1600 train_time:58283ms step_avg:51.08ms
step:1142/1600 train_time:58369ms step_avg:51.11ms
step:1143/1600 train_time:58458ms step_avg:51.14ms
step:1144/1600 train_time:58542ms step_avg:51.17ms
step:1145/1600 train_time:58631ms step_avg:51.21ms
step:1146/1600 train_time:58716ms step_avg:51.24ms
step:1147/1600 train_time:58804ms step_avg:51.27ms
step:1148/1600 train_time:58889ms step_avg:51.30ms
step:1149/1600 train_time:58977ms step_avg:51.33ms
step:1150/1600 train_time:59062ms step_avg:51.36ms
step:1151/1600 train_time:59151ms step_avg:51.39ms
step:1152/1600 train_time:59239ms step_avg:51.42ms
step:1153/1600 train_time:59325ms step_avg:51.45ms
step:1154/1600 train_time:59414ms step_avg:51.49ms
step:1155/1600 train_time:59501ms step_avg:51.52ms
step:1156/1600 train_time:59586ms step_avg:51.54ms
step:1157/1600 train_time:59673ms step_avg:51.58ms
step:1158/1600 train_time:59758ms step_avg:51.60ms
step:1159/1600 train_time:59846ms step_avg:51.64ms
step:1160/1600 train_time:59932ms step_avg:51.67ms
step:1161/1600 train_time:60021ms step_avg:51.70ms
step:1162/1600 train_time:60105ms step_avg:51.73ms
step:1163/1600 train_time:60193ms step_avg:51.76ms
step:1164/1600 train_time:60280ms step_avg:51.79ms
step:1165/1600 train_time:60367ms step_avg:51.82ms
step:1166/1600 train_time:60454ms step_avg:51.85ms
step:1167/1600 train_time:60541ms step_avg:51.88ms
step:1168/1600 train_time:60626ms step_avg:51.91ms
step:1169/1600 train_time:60714ms step_avg:51.94ms
step:1170/1600 train_time:60799ms step_avg:51.96ms
step:1171/1600 train_time:60888ms step_avg:52.00ms
step:1172/1600 train_time:60973ms step_avg:52.03ms
step:1173/1600 train_time:61062ms step_avg:52.06ms
step:1174/1600 train_time:61147ms step_avg:52.08ms
step:1175/1600 train_time:61235ms step_avg:52.11ms
step:1176/1600 train_time:61321ms step_avg:52.14ms
step:1177/1600 train_time:61408ms step_avg:52.17ms
step:1178/1600 train_time:61493ms step_avg:52.20ms
step:1179/1600 train_time:61582ms step_avg:52.23ms
step:1180/1600 train_time:61668ms step_avg:52.26ms
step:1181/1600 train_time:61756ms step_avg:52.29ms
step:1182/1600 train_time:61841ms step_avg:52.32ms
step:1183/1600 train_time:61929ms step_avg:52.35ms
step:1184/1600 train_time:62014ms step_avg:52.38ms
step:1185/1600 train_time:62103ms step_avg:52.41ms
step:1186/1600 train_time:62187ms step_avg:52.43ms
step:1187/1600 train_time:62276ms step_avg:52.47ms
step:1188/1600 train_time:62361ms step_avg:52.49ms
step:1189/1600 train_time:62450ms step_avg:52.52ms
step:1190/1600 train_time:62535ms step_avg:52.55ms
step:1191/1600 train_time:62624ms step_avg:52.58ms
step:1192/1600 train_time:62709ms step_avg:52.61ms
step:1193/1600 train_time:62798ms step_avg:52.64ms
step:1194/1600 train_time:62883ms step_avg:52.67ms
step:1195/1600 train_time:62971ms step_avg:52.70ms
step:1196/1600 train_time:63056ms step_avg:52.72ms
step:1197/1600 train_time:63145ms step_avg:52.75ms
step:1198/1600 train_time:63231ms step_avg:52.78ms
step:1199/1600 train_time:63319ms step_avg:52.81ms
step:1200/1600 train_time:63405ms step_avg:52.84ms
step:1201/1600 train_time:63493ms step_avg:52.87ms
step:1202/1600 train_time:63578ms step_avg:52.89ms
step:1203/1600 train_time:63666ms step_avg:52.92ms
step:1204/1600 train_time:63751ms step_avg:52.95ms
step:1205/1600 train_time:63839ms step_avg:52.98ms
step:1206/1600 train_time:63925ms step_avg:53.01ms
step:1207/1600 train_time:64013ms step_avg:53.04ms
step:1208/1600 train_time:64099ms step_avg:53.06ms
step:1209/1600 train_time:64187ms step_avg:53.09ms
step:1210/1600 train_time:64273ms step_avg:53.12ms
step:1211/1600 train_time:64361ms step_avg:53.15ms
step:1212/1600 train_time:64447ms step_avg:53.17ms
step:1213/1600 train_time:64537ms step_avg:53.20ms
step:1214/1600 train_time:64621ms step_avg:53.23ms
step:1215/1600 train_time:64709ms step_avg:53.26ms
step:1216/1600 train_time:64794ms step_avg:53.28ms
step:1217/1600 train_time:64882ms step_avg:53.31ms
step:1218/1600 train_time:64968ms step_avg:53.34ms
step:1219/1600 train_time:65056ms step_avg:53.37ms
step:1220/1600 train_time:65140ms step_avg:53.39ms
step:1221/1600 train_time:65228ms step_avg:53.42ms
step:1222/1600 train_time:65313ms step_avg:53.45ms
step:1223/1600 train_time:65402ms step_avg:53.48ms
step:1224/1600 train_time:65487ms step_avg:53.50ms
step:1225/1600 train_time:65575ms step_avg:53.53ms
step:1226/1600 train_time:65660ms step_avg:53.56ms
step:1227/1600 train_time:65749ms step_avg:53.59ms
step:1228/1600 train_time:65834ms step_avg:53.61ms
step:1229/1600 train_time:65925ms step_avg:53.64ms
step:1230/1600 train_time:66008ms step_avg:53.67ms
step:1231/1600 train_time:66096ms step_avg:53.69ms
step:1232/1600 train_time:66181ms step_avg:53.72ms
step:1233/1600 train_time:66269ms step_avg:53.75ms
step:1234/1600 train_time:66354ms step_avg:53.77ms
step:1235/1600 train_time:66443ms step_avg:53.80ms
step:1236/1600 train_time:66528ms step_avg:53.83ms
step:1237/1600 train_time:66616ms step_avg:53.85ms
step:1238/1600 train_time:66701ms step_avg:53.88ms
step:1239/1600 train_time:66790ms step_avg:53.91ms
step:1240/1600 train_time:66875ms step_avg:53.93ms
step:1241/1600 train_time:66964ms step_avg:53.96ms
step:1242/1600 train_time:67048ms step_avg:53.98ms
step:1243/1600 train_time:67138ms step_avg:54.01ms
step:1244/1600 train_time:67223ms step_avg:54.04ms
step:1245/1600 train_time:67311ms step_avg:54.07ms
step:1246/1600 train_time:67397ms step_avg:54.09ms
step:1247/1600 train_time:67485ms step_avg:54.12ms
step:1248/1600 train_time:67570ms step_avg:54.14ms
step:1249/1600 train_time:67658ms step_avg:54.17ms
step:1250/1600 train_time:67744ms step_avg:54.20ms
step:1250/1600 val_loss:3.4172 train_time:67818ms step_avg:54.25ms
step:1251/1600 train_time:67837ms step_avg:54.23ms
step:1252/1600 train_time:67923ms step_avg:54.25ms
step:1253/1600 train_time:68015ms step_avg:54.28ms
step:1254/1600 train_time:68100ms step_avg:54.31ms
step:1255/1600 train_time:68189ms step_avg:54.33ms
step:1256/1600 train_time:68274ms step_avg:54.36ms
step:1257/1600 train_time:68361ms step_avg:54.38ms
step:1258/1600 train_time:68446ms step_avg:54.41ms
step:1259/1600 train_time:68535ms step_avg:54.44ms
step:1260/1600 train_time:68617ms step_avg:54.46ms
step:1261/1600 train_time:68703ms step_avg:54.48ms
step:1262/1600 train_time:68789ms step_avg:54.51ms
step:1263/1600 train_time:68878ms step_avg:54.54ms
step:1264/1600 train_time:68966ms step_avg:54.56ms
step:1265/1600 train_time:69055ms step_avg:54.59ms
step:1266/1600 train_time:69141ms step_avg:54.61ms
step:1267/1600 train_time:69230ms step_avg:54.64ms
step:1268/1600 train_time:69314ms step_avg:54.66ms
step:1269/1600 train_time:69403ms step_avg:54.69ms
step:1270/1600 train_time:69487ms step_avg:54.71ms
step:1271/1600 train_time:69575ms step_avg:54.74ms
step:1272/1600 train_time:69658ms step_avg:54.76ms
step:1273/1600 train_time:69747ms step_avg:54.79ms
step:1274/1600 train_time:69832ms step_avg:54.81ms
step:1275/1600 train_time:69923ms step_avg:54.84ms
step:1276/1600 train_time:70010ms step_avg:54.87ms
step:1277/1600 train_time:70098ms step_avg:54.89ms
step:1278/1600 train_time:70184ms step_avg:54.92ms
step:1279/1600 train_time:70272ms step_avg:54.94ms
step:1280/1600 train_time:70357ms step_avg:54.97ms
step:1281/1600 train_time:70445ms step_avg:54.99ms
step:1282/1600 train_time:70529ms step_avg:55.01ms
step:1283/1600 train_time:70616ms step_avg:55.04ms
step:1284/1600 train_time:70701ms step_avg:55.06ms
step:1285/1600 train_time:70790ms step_avg:55.09ms
step:1286/1600 train_time:70875ms step_avg:55.11ms
step:1287/1600 train_time:70965ms step_avg:55.14ms
step:1288/1600 train_time:71051ms step_avg:55.16ms
step:1289/1600 train_time:71139ms step_avg:55.19ms
step:1290/1600 train_time:71224ms step_avg:55.21ms
step:1291/1600 train_time:71313ms step_avg:55.24ms
step:1292/1600 train_time:71397ms step_avg:55.26ms
step:1293/1600 train_time:71484ms step_avg:55.29ms
step:1294/1600 train_time:71569ms step_avg:55.31ms
step:1295/1600 train_time:71658ms step_avg:55.33ms
step:1296/1600 train_time:71742ms step_avg:55.36ms
step:1297/1600 train_time:71831ms step_avg:55.38ms
step:1298/1600 train_time:71916ms step_avg:55.41ms
step:1299/1600 train_time:72005ms step_avg:55.43ms
step:1300/1600 train_time:72091ms step_avg:55.45ms
step:1301/1600 train_time:72179ms step_avg:55.48ms
step:1302/1600 train_time:72264ms step_avg:55.50ms
step:1303/1600 train_time:72352ms step_avg:55.53ms
step:1304/1600 train_time:72437ms step_avg:55.55ms
step:1305/1600 train_time:72524ms step_avg:55.57ms
step:1306/1600 train_time:72609ms step_avg:55.60ms
step:1307/1600 train_time:72697ms step_avg:55.62ms
step:1308/1600 train_time:72782ms step_avg:55.64ms
step:1309/1600 train_time:72871ms step_avg:55.67ms
step:1310/1600 train_time:72956ms step_avg:55.69ms
step:1311/1600 train_time:73045ms step_avg:55.72ms
step:1312/1600 train_time:73132ms step_avg:55.74ms
step:1313/1600 train_time:73220ms step_avg:55.77ms
step:1314/1600 train_time:73306ms step_avg:55.79ms
step:1315/1600 train_time:73395ms step_avg:55.81ms
step:1316/1600 train_time:73480ms step_avg:55.84ms
step:1317/1600 train_time:73567ms step_avg:55.86ms
step:1318/1600 train_time:73652ms step_avg:55.88ms
step:1319/1600 train_time:73740ms step_avg:55.91ms
step:1320/1600 train_time:73825ms step_avg:55.93ms
step:1321/1600 train_time:73913ms step_avg:55.95ms
step:1322/1600 train_time:73998ms step_avg:55.97ms
step:1323/1600 train_time:74087ms step_avg:56.00ms
step:1324/1600 train_time:74173ms step_avg:56.02ms
step:1325/1600 train_time:74263ms step_avg:56.05ms
step:1326/1600 train_time:74348ms step_avg:56.07ms
step:1327/1600 train_time:74437ms step_avg:56.09ms
step:1328/1600 train_time:74522ms step_avg:56.12ms
step:1329/1600 train_time:74609ms step_avg:56.14ms
step:1330/1600 train_time:74694ms step_avg:56.16ms
step:1331/1600 train_time:74783ms step_avg:56.19ms
step:1332/1600 train_time:74868ms step_avg:56.21ms
step:1333/1600 train_time:74956ms step_avg:56.23ms
step:1334/1600 train_time:75041ms step_avg:56.25ms
step:1335/1600 train_time:75129ms step_avg:56.28ms
step:1336/1600 train_time:75214ms step_avg:56.30ms
step:1337/1600 train_time:75302ms step_avg:56.32ms
step:1338/1600 train_time:75388ms step_avg:56.34ms
step:1339/1600 train_time:75476ms step_avg:56.37ms
step:1340/1600 train_time:75561ms step_avg:56.39ms
step:1341/1600 train_time:75649ms step_avg:56.41ms
step:1342/1600 train_time:75734ms step_avg:56.43ms
step:1343/1600 train_time:75823ms step_avg:56.46ms
step:1344/1600 train_time:75908ms step_avg:56.48ms
step:1345/1600 train_time:75997ms step_avg:56.50ms
step:1346/1600 train_time:76082ms step_avg:56.52ms
step:1347/1600 train_time:76170ms step_avg:56.55ms
step:1348/1600 train_time:76255ms step_avg:56.57ms
step:1349/1600 train_time:76343ms step_avg:56.59ms
step:1350/1600 train_time:76428ms step_avg:56.61ms
step:1351/1600 train_time:76517ms step_avg:56.64ms
step:1352/1600 train_time:76601ms step_avg:56.66ms
step:1353/1600 train_time:76689ms step_avg:56.68ms
step:1354/1600 train_time:76774ms step_avg:56.70ms
step:1355/1600 train_time:76862ms step_avg:56.72ms
step:1356/1600 train_time:76949ms step_avg:56.75ms
step:1357/1600 train_time:77036ms step_avg:56.77ms
step:1358/1600 train_time:77122ms step_avg:56.79ms
step:1359/1600 train_time:77210ms step_avg:56.81ms
step:1360/1600 train_time:77296ms step_avg:56.84ms
step:1361/1600 train_time:77384ms step_avg:56.86ms
step:1362/1600 train_time:77469ms step_avg:56.88ms
step:1363/1600 train_time:77557ms step_avg:56.90ms
step:1364/1600 train_time:77642ms step_avg:56.92ms
step:1365/1600 train_time:77731ms step_avg:56.95ms
step:1366/1600 train_time:77820ms step_avg:56.97ms
step:1367/1600 train_time:77906ms step_avg:56.99ms
step:1368/1600 train_time:77990ms step_avg:57.01ms
step:1369/1600 train_time:78077ms step_avg:57.03ms
step:1370/1600 train_time:78162ms step_avg:57.05ms
step:1371/1600 train_time:78250ms step_avg:57.08ms
step:1372/1600 train_time:78337ms step_avg:57.10ms
step:1373/1600 train_time:78424ms step_avg:57.12ms
step:1374/1600 train_time:78509ms step_avg:57.14ms
step:1375/1600 train_time:78598ms step_avg:57.16ms
step:1376/1600 train_time:78683ms step_avg:57.18ms
step:1377/1600 train_time:78772ms step_avg:57.21ms
step:1378/1600 train_time:78858ms step_avg:57.23ms
step:1379/1600 train_time:78946ms step_avg:57.25ms
step:1380/1600 train_time:79032ms step_avg:57.27ms
step:1381/1600 train_time:79120ms step_avg:57.29ms
step:1382/1600 train_time:79205ms step_avg:57.31ms
step:1383/1600 train_time:79293ms step_avg:57.33ms
step:1384/1600 train_time:79378ms step_avg:57.35ms
step:1385/1600 train_time:79467ms step_avg:57.38ms
step:1386/1600 train_time:79551ms step_avg:57.40ms
step:1387/1600 train_time:79640ms step_avg:57.42ms
step:1388/1600 train_time:79724ms step_avg:57.44ms
step:1389/1600 train_time:79814ms step_avg:57.46ms
step:1390/1600 train_time:79899ms step_avg:57.48ms
step:1391/1600 train_time:79987ms step_avg:57.50ms
step:1392/1600 train_time:80072ms step_avg:57.52ms
step:1393/1600 train_time:80161ms step_avg:57.55ms
step:1394/1600 train_time:80246ms step_avg:57.57ms
step:1395/1600 train_time:80334ms step_avg:57.59ms
step:1396/1600 train_time:80421ms step_avg:57.61ms
step:1397/1600 train_time:80507ms step_avg:57.63ms
step:1398/1600 train_time:80593ms step_avg:57.65ms
step:1399/1600 train_time:80681ms step_avg:57.67ms
step:1400/1600 train_time:80766ms step_avg:57.69ms
step:1401/1600 train_time:80855ms step_avg:57.71ms
step:1402/1600 train_time:80940ms step_avg:57.73ms
step:1403/1600 train_time:81028ms step_avg:57.75ms
step:1404/1600 train_time:81112ms step_avg:57.77ms
step:1405/1600 train_time:81201ms step_avg:57.79ms
step:1406/1600 train_time:81286ms step_avg:57.81ms
step:1407/1600 train_time:81375ms step_avg:57.84ms
step:1408/1600 train_time:81461ms step_avg:57.86ms
step:1409/1600 train_time:81549ms step_avg:57.88ms
step:1410/1600 train_time:81635ms step_avg:57.90ms
step:1411/1600 train_time:81724ms step_avg:57.92ms
step:1412/1600 train_time:81809ms step_avg:57.94ms
step:1413/1600 train_time:81897ms step_avg:57.96ms
step:1414/1600 train_time:81982ms step_avg:57.98ms
step:1415/1600 train_time:82070ms step_avg:58.00ms
step:1416/1600 train_time:82155ms step_avg:58.02ms
step:1417/1600 train_time:82244ms step_avg:58.04ms
step:1418/1600 train_time:82329ms step_avg:58.06ms
step:1419/1600 train_time:82417ms step_avg:58.08ms
step:1420/1600 train_time:82502ms step_avg:58.10ms
step:1421/1600 train_time:82590ms step_avg:58.12ms
step:1422/1600 train_time:82676ms step_avg:58.14ms
step:1423/1600 train_time:82765ms step_avg:58.16ms
step:1424/1600 train_time:82850ms step_avg:58.18ms
step:1425/1600 train_time:82938ms step_avg:58.20ms
step:1426/1600 train_time:83023ms step_avg:58.22ms
step:1427/1600 train_time:83111ms step_avg:58.24ms
step:1428/1600 train_time:83197ms step_avg:58.26ms
step:1429/1600 train_time:83285ms step_avg:58.28ms
step:1430/1600 train_time:83371ms step_avg:58.30ms
step:1431/1600 train_time:83463ms step_avg:58.32ms
step:1432/1600 train_time:83549ms step_avg:58.34ms
step:1433/1600 train_time:83636ms step_avg:58.36ms
step:1434/1600 train_time:83720ms step_avg:58.38ms
step:1435/1600 train_time:83808ms step_avg:58.40ms
step:1436/1600 train_time:83893ms step_avg:58.42ms
step:1437/1600 train_time:83981ms step_avg:58.44ms
step:1438/1600 train_time:84067ms step_avg:58.46ms
step:1439/1600 train_time:84155ms step_avg:58.48ms
step:1440/1600 train_time:84240ms step_avg:58.50ms
step:1441/1600 train_time:84328ms step_avg:58.52ms
step:1442/1600 train_time:84412ms step_avg:58.54ms
step:1443/1600 train_time:84502ms step_avg:58.56ms
step:1444/1600 train_time:84587ms step_avg:58.58ms
step:1445/1600 train_time:84676ms step_avg:58.60ms
step:1446/1600 train_time:84760ms step_avg:58.62ms
step:1447/1600 train_time:84848ms step_avg:58.64ms
step:1448/1600 train_time:84933ms step_avg:58.66ms
step:1449/1600 train_time:85021ms step_avg:58.68ms
step:1450/1600 train_time:85106ms step_avg:58.69ms
step:1451/1600 train_time:85194ms step_avg:58.71ms
step:1452/1600 train_time:85280ms step_avg:58.73ms
step:1453/1600 train_time:85367ms step_avg:58.75ms
step:1454/1600 train_time:85452ms step_avg:58.77ms
step:1455/1600 train_time:85541ms step_avg:58.79ms
step:1456/1600 train_time:85626ms step_avg:58.81ms
step:1457/1600 train_time:85715ms step_avg:58.83ms
step:1458/1600 train_time:85800ms step_avg:58.85ms
step:1459/1600 train_time:85888ms step_avg:58.87ms
step:1460/1600 train_time:85973ms step_avg:58.89ms
step:1461/1600 train_time:86062ms step_avg:58.91ms
step:1462/1600 train_time:86148ms step_avg:58.92ms
step:1463/1600 train_time:86236ms step_avg:58.94ms
step:1464/1600 train_time:86322ms step_avg:58.96ms
step:1465/1600 train_time:86410ms step_avg:58.98ms
step:1466/1600 train_time:86495ms step_avg:59.00ms
step:1467/1600 train_time:86585ms step_avg:59.02ms
step:1468/1600 train_time:86670ms step_avg:59.04ms
step:1469/1600 train_time:86758ms step_avg:59.06ms
step:1470/1600 train_time:86843ms step_avg:59.08ms
step:1471/1600 train_time:86932ms step_avg:59.10ms
step:1472/1600 train_time:87017ms step_avg:59.11ms
step:1473/1600 train_time:87105ms step_avg:59.13ms
step:1474/1600 train_time:87190ms step_avg:59.15ms
step:1475/1600 train_time:87279ms step_avg:59.17ms
step:1476/1600 train_time:87363ms step_avg:59.19ms
step:1477/1600 train_time:87452ms step_avg:59.21ms
step:1478/1600 train_time:87537ms step_avg:59.23ms
step:1479/1600 train_time:87626ms step_avg:59.25ms
step:1480/1600 train_time:87712ms step_avg:59.26ms
step:1481/1600 train_time:87801ms step_avg:59.28ms
step:1482/1600 train_time:87885ms step_avg:59.30ms
step:1483/1600 train_time:87974ms step_avg:59.32ms
step:1484/1600 train_time:88059ms step_avg:59.34ms
step:1485/1600 train_time:88147ms step_avg:59.36ms
step:1486/1600 train_time:88231ms step_avg:59.38ms
step:1487/1600 train_time:88320ms step_avg:59.39ms
step:1488/1600 train_time:88407ms step_avg:59.41ms
step:1489/1600 train_time:88494ms step_avg:59.43ms
step:1490/1600 train_time:88579ms step_avg:59.45ms
step:1491/1600 train_time:88667ms step_avg:59.47ms
step:1492/1600 train_time:88753ms step_avg:59.49ms
step:1493/1600 train_time:88840ms step_avg:59.50ms
step:1494/1600 train_time:88925ms step_avg:59.52ms
step:1495/1600 train_time:89014ms step_avg:59.54ms
step:1496/1600 train_time:89100ms step_avg:59.56ms
step:1497/1600 train_time:89188ms step_avg:59.58ms
step:1498/1600 train_time:89273ms step_avg:59.59ms
step:1499/1600 train_time:89362ms step_avg:59.61ms
step:1500/1600 train_time:89447ms step_avg:59.63ms
step:1500/1600 val_loss:3.3084 train_time:89519ms step_avg:59.68ms
step:1501/1600 train_time:89539ms step_avg:59.65ms
step:1502/1600 train_time:89624ms step_avg:59.67ms
step:1503/1600 train_time:89716ms step_avg:59.69ms
step:1504/1600 train_time:89801ms step_avg:59.71ms
step:1505/1600 train_time:89889ms step_avg:59.73ms
step:1506/1600 train_time:89973ms step_avg:59.74ms
step:1507/1600 train_time:90060ms step_avg:59.76ms
step:1508/1600 train_time:90145ms step_avg:59.78ms
step:1509/1600 train_time:90232ms step_avg:59.80ms
step:1510/1600 train_time:90317ms step_avg:59.81ms
step:1511/1600 train_time:90405ms step_avg:59.83ms
step:1512/1600 train_time:90491ms step_avg:59.85ms
step:1513/1600 train_time:90582ms step_avg:59.87ms
step:1514/1600 train_time:90669ms step_avg:59.89ms
step:1515/1600 train_time:90759ms step_avg:59.91ms
step:1516/1600 train_time:90844ms step_avg:59.92ms
step:1517/1600 train_time:90932ms step_avg:59.94ms
step:1518/1600 train_time:91017ms step_avg:59.96ms
step:1519/1600 train_time:91104ms step_avg:59.98ms
step:1520/1600 train_time:91190ms step_avg:59.99ms
step:1521/1600 train_time:91277ms step_avg:60.01ms
step:1522/1600 train_time:91362ms step_avg:60.03ms
step:1523/1600 train_time:91450ms step_avg:60.05ms
step:1524/1600 train_time:91537ms step_avg:60.06ms
step:1525/1600 train_time:91627ms step_avg:60.08ms
step:1526/1600 train_time:91713ms step_avg:60.10ms
step:1527/1600 train_time:91801ms step_avg:60.12ms
step:1528/1600 train_time:91887ms step_avg:60.14ms
step:1529/1600 train_time:91976ms step_avg:60.15ms
step:1530/1600 train_time:92060ms step_avg:60.17ms
step:1531/1600 train_time:92147ms step_avg:60.19ms
step:1532/1600 train_time:92233ms step_avg:60.20ms
step:1533/1600 train_time:92320ms step_avg:60.22ms
step:1534/1600 train_time:92405ms step_avg:60.24ms
step:1535/1600 train_time:92494ms step_avg:60.26ms
step:1536/1600 train_time:92581ms step_avg:60.27ms
step:1537/1600 train_time:92670ms step_avg:60.29ms
step:1538/1600 train_time:92757ms step_avg:60.31ms
step:1539/1600 train_time:92844ms step_avg:60.33ms
step:1540/1600 train_time:92929ms step_avg:60.34ms
step:1541/1600 train_time:93017ms step_avg:60.36ms
step:1542/1600 train_time:93102ms step_avg:60.38ms
step:1543/1600 train_time:93189ms step_avg:60.39ms
step:1544/1600 train_time:93273ms step_avg:60.41ms
step:1545/1600 train_time:93361ms step_avg:60.43ms
step:1546/1600 train_time:93446ms step_avg:60.44ms
step:1547/1600 train_time:93536ms step_avg:60.46ms
step:1548/1600 train_time:93622ms step_avg:60.48ms
step:1549/1600 train_time:93711ms step_avg:60.50ms
step:1550/1600 train_time:93796ms step_avg:60.51ms
step:1551/1600 train_time:93885ms step_avg:60.53ms
step:1552/1600 train_time:93970ms step_avg:60.55ms
step:1553/1600 train_time:94058ms step_avg:60.57ms
step:1554/1600 train_time:94143ms step_avg:60.58ms
step:1555/1600 train_time:94230ms step_avg:60.60ms
step:1556/1600 train_time:94315ms step_avg:60.61ms
step:1557/1600 train_time:94404ms step_avg:60.63ms
step:1558/1600 train_time:94489ms step_avg:60.65ms
step:1559/1600 train_time:94578ms step_avg:60.67ms
step:1560/1600 train_time:94663ms step_avg:60.68ms
step:1561/1600 train_time:94759ms step_avg:60.70ms
step:1562/1600 train_time:94844ms step_avg:60.72ms
step:1563/1600 train_time:94932ms step_avg:60.74ms
step:1564/1600 train_time:95018ms step_avg:60.75ms
step:1565/1600 train_time:95106ms step_avg:60.77ms
step:1566/1600 train_time:95191ms step_avg:60.79ms
step:1567/1600 train_time:95279ms step_avg:60.80ms
step:1568/1600 train_time:95365ms step_avg:60.82ms
step:1569/1600 train_time:95453ms step_avg:60.84ms
step:1570/1600 train_time:95538ms step_avg:60.85ms
step:1571/1600 train_time:95627ms step_avg:60.87ms
step:1572/1600 train_time:95714ms step_avg:60.89ms
step:1573/1600 train_time:95802ms step_avg:60.90ms
step:1574/1600 train_time:95889ms step_avg:60.92ms
step:1575/1600 train_time:95978ms step_avg:60.94ms
step:1576/1600 train_time:96063ms step_avg:60.95ms
step:1577/1600 train_time:96154ms step_avg:60.97ms
step:1578/1600 train_time:96238ms step_avg:60.99ms
step:1579/1600 train_time:96327ms step_avg:61.01ms
step:1580/1600 train_time:96412ms step_avg:61.02ms
step:1581/1600 train_time:96500ms step_avg:61.04ms
step:1582/1600 train_time:96586ms step_avg:61.05ms
step:1583/1600 train_time:96676ms step_avg:61.07ms
step:1584/1600 train_time:96761ms step_avg:61.09ms
step:1585/1600 train_time:96850ms step_avg:61.10ms
step:1586/1600 train_time:96935ms step_avg:61.12ms
step:1587/1600 train_time:97024ms step_avg:61.14ms
step:1588/1600 train_time:97108ms step_avg:61.15ms
step:1589/1600 train_time:97196ms step_avg:61.17ms
step:1590/1600 train_time:97282ms step_avg:61.18ms
step:1591/1600 train_time:97370ms step_avg:61.20ms
step:1592/1600 train_time:97456ms step_avg:61.22ms
step:1593/1600 train_time:97545ms step_avg:61.23ms
step:1594/1600 train_time:97632ms step_avg:61.25ms
step:1595/1600 train_time:97719ms step_avg:61.27ms
step:1596/1600 train_time:97804ms step_avg:61.28ms
step:1597/1600 train_time:97893ms step_avg:61.30ms
step:1598/1600 train_time:97978ms step_avg:61.31ms
step:1599/1600 train_time:98067ms step_avg:61.33ms
step:1600/1600 train_time:98152ms step_avg:61.34ms
step:1600/1600 val_loss:3.2791 train_time:98225ms step_avg:61.39ms
peak memory allocated: 30264 MiB reserved: 46238 MiB