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, pairwise
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

# -----------------------------------------------------------------------------
# Distributed training setup
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
grad_scale = 2 / grad_accum_steps # consistent grad magnitudes between different num_devices
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.

# -----------------------------------------------------------------------------
# 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, paired=False):
        super().__init__()
        self.head_dim = head_dim
        self.max_seq_len = max_seq_len
        self.paired = paired
        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)
        if not self.paired:
            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
            )
        else:
            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 = 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)
        if not self.paired:
            theta = torch.outer(t, self.angular_freq)
            self.factor1.copy_(theta.cos())
            self.factor2.copy_(theta.sin())
        else:
            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, paired: bool = False):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.dim = dim
        self.hdim = num_heads * head_dim
        self.paired = paired
        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)
        max_len = args.train_max_seq_len if self.training else (args.val_batch_size // (grad_accum_steps * world_size))

        q, k = norm(q), norm(k) # QK norm @Grad62304977

        if not self.paired:
            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

        else:
            # Paired heads: adjacent heads' queries attend to each other's keys.
            # Two copies of the input stream are interleaved to achieve this, which:
            # - doubles the length of each sequence
            # - halves the effective window size
            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)

            seqlens = 2 * seqlens
            max_len = 2 * max_len

        # 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 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
        self.attn = CausalSelfAttention(dim, head_dim, num_heads, paired=use_paired_head) if has_attn else 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
        self.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.Parameter(torch.zeros(5 * self.vocab_size, model_dim, dtype=torch.bfloat16))
        self.value_embeds.label = 'value_embed'

        # parameter banks for attention and value embedding gate weights
        self.attn_gate_bank = nn.Parameter(torch.zeros(10, num_heads, 12)) # 10 layers
        self.attn_gate_bank.label = 'attn_gate_bank'
        self.ve_gate_bank = nn.Parameter(torch.zeros(5, num_heads, 12)) # 5 unique gates
        self.ve_gate_bank.label = 've_gate_bank'

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

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

        hdim = num_heads * head_dim
        mlp_hdim = 4 * model_dim

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

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

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

        # improved init scale by @YouJiacheng and @srashedll
        std = 0.5 * model_dim ** -0.5
        bound = (3 ** 0.5) * std
        with torch.no_grad():
            self.attn_bank.uniform_(-bound, bound)
            self.mlp_bank[:, 0, :, :].uniform_(-bound, 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 = Yarn(head_dim, max_seq_len, paired=True)
        # 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, self.vocab_size, use_fp8=use_fp8, x_s=100/448, w_s=1.6/448, grad_s=grad_scale * 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(self.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'

    @staticmethod
    @torch.compile(dynamic=False, fullgraph=True)
    def _compute_bigram_hash(x: Tensor, mod: int) -> Tensor:
        """
        Computes bigram hash on GPU for each position using [prev_token, curr_token].
        Mathematically identical to the CPU version but computed on device.
        """
        rand_int_1 = 36313
        rand_int_2 = 27191
        result = torch.empty_like(x)
        result[0] = mod
        result[1:] = torch.bitwise_xor(rand_int_1 * x[1:], rand_int_2 * x[:-1]) % mod
        return result

    def forward(self, input_seq: Tensor, target_seq: Tensor, seqlens: 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
        bm_sizes = [ws_short, ws_short, ws_short, ws_long, ws_short, ws_short, None, ws_short, ws_short, ws_short, ws_long]
        assert len(bm_sizes) == self.num_layers
        key_offset = [b==ws_long 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)
        # Compute bigram hash on GPU (moved from CPU data loader)
        bigram_seq = self._compute_bigram_hash(input_seq, args.bigram_vocab_size - 1)
        x0_bigram = self.bigram_embed(bigram_seq)[None]

        # Value embeddings - always computed (not precomputed)
        ve = self.value_embeds.view(5, self.vocab_size, -1)[:, input_seq]
        # 01 ... 234 structure on token value embeddings by @photomz
        ve = [ve[0], ve[1]] + [None] * (self.num_layers - 5) + [ve[2], ve[3], ve[4]]
        assert len(ve) == self.num_layers

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

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

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

        for i in range(self.num_layers):
            yarn = self.yarn_paired_head if i in self.paired_head_layers else self.yarn
            attn_args = AttnArgs(
                ve=ve[i],
                sa_lambdas=sa_lambdas[i],
                seqlens=seqlens,
                bm_size=bm_sizes[i],
                yarn=yarn,
                key_offset=key_offset[i],
                attn_gate_w=attn_gates[i],
                ve_gate_w=ve_gates[i]
            )
            if i in skip_out:
                skip_gate_out = torch.sigmoid(skip_lambda) * 2 * torch.sigmoid(self.skip_gate(x0[..., :self.skip_gate.weight.size(-1)]))
                x = x + skip_gate_out * skip_connections.pop()
            if i == 0:
                x = (resid_lambdas[0] + x0_lambdas[0]) * x + bigram_lambdas[0] * x0_bigram
            else:
                x = resid_lambdas[i] * x + x0_lambdas[i] * x0 + bigram_lambdas[i] * x0_bigram

            # Get weights for this layer from banks
            qkvo_w = attn_weights[self.layer_to_attn_idx[i]] if i in self.layer_to_attn_idx else None
            c_fc = mlp_fcs[self.layer_to_mlp_idx[i]] if i in self.layer_to_mlp_idx else None
            c_proj = mlp_projs[self.layer_to_mlp_idx[i]] if i in self.layer_to_mlp_idx else None

            x = self.blocks[i](x, attn_args, qkvo_w, c_fc, c_proj)
            if i in skip_in:
                skip_connections.append(x)
            if i == backout_layer:
                x_backout = x

        # back out contributions from first 7 layers that are only required for downstream context and not direct prediction
        x -= backout_lambda * x_backout
        x = norm(x)
        logits = self.lm_head(x)
        # @Grad62304977 added tanh softcapping following Gemma 2 paper, @KoszarskyB reduced it from 30 to 15
        # @YouJiacheng shifted it by +15 (2*sigmoid(2*x)=tanh(x)+1). @classiclarryd updated to 23*sigmoid((logits+5)/7.5)
        if self.training:
            losses = FusedSoftcappedCrossEntropy.apply(logits.view(-1, logits.size(-1)), target_seq, mtp_weights, 23.0, 5.0, 7.5)
            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 Shard:
    def __init__(self, tokens: Tensor, world_size: int = 1):
        self.tokens = tokens
        self.size = tokens.numel()
        self.world_size = world_size
        self.i = 0

        # Partial index now, full index async
        self.bos_idx = (tokens[:6_000_000] == BOS_ID).nonzero(as_tuple=True)[0].to(torch.int64).cpu().numpy()
        self._full_idx = None
        self._loader_thread = None
        self._ready = threading.Event()
        self._loader_thread = threading.Thread(target=self._scan)
        self._loader_thread.start()

    def _scan(self):
        self._full_idx = (self.tokens == BOS_ID).nonzero(as_tuple=True)[0].to(torch.int64).cpu().numpy()
        self._ready.set()

    def _maybe_switch(self):
        # Switch to full index as soon as async scan completes
        if self.bos_idx is not self._full_idx and self._ready.is_set():
            self._loader_thread.join()
            self.bos_idx = self._full_idx

    def next_batch(self, num_tokens_local: int, max_seq_len: int):
        self._maybe_switch()
        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
        return starts, ends

    @staticmethod
    def load_async(file: Path, world_size: int = 1):
        """Returns getter function for async shard loading"""
        result = {}
        ready = threading.Event()
        def load():
            tokens = _load_data_shard(file)
            result['shard'] = Shard(tokens, world_size)
            ready.set()
        thread = threading.Thread(target=load)
        thread.start()
        def get():
            ready.wait()
            thread.join()
            return result['shard']
        return get

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:
        shard = Shard(tokens, world_size)
        next_shard_getter = Shard.load_async(next(file_iter), world_size)
    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 = shard.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.
                shard = next_shard_getter()
                tokens = shard.tokens
                try:
                    next_shard_getter = Shard.load_async(next(file_iter), world_size)
                except StopIteration:
                    next_shard_getter = None  # no more shards to preload
                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 hash computation moved to GPU in forward()

        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),
        )

        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

@dataclass
class Hyperparameters:
    # data
    data_path = os.environ.get("DATA_PATH", ".")
    train_files: str = os.path.join(data_path, "data/fineweb10B/fineweb_train_*.bin") # input .bin to train on
    val_files: str = os.path.join(data_path, "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_max_seq_len: int = 128 * 16
    val_batch_size: int = 4 * 64 * 1024 * 8
    # schedule
    num_scheduled_iterations: int = 1515  # 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
    # 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
    # bigram hash embedding
    bigram_vocab_size: int = 50304 * 5

args = Hyperparameters()

@dataclass
class TrainingStage:
    lr_mul: float
    batch_size: int
    window_sizes: tuple[int, int]  # (short, long) in block units
    mtp_weights_start: list[float]
    mtp_weights_end: list[float]
    duration: float = None

class TrainingSchedule:
    """
    Training schedule initialized via TRAINING_STAGES
        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
        5. Batch size schedule of 8 -> 16 -> 24
        6. Post training extension of long windows from 13 to 20
    """

    def __init__(self, stages: list[TrainingStage], scheduled_iterations: int, extension_iterations: int,
                 cooldown_frac: float = 0.5, split_embed_stage: int = 2, ws_post_yarn_ext: int = 20):
        self.stages = stages
        self.scheduled_iterations = scheduled_iterations
        self.cooldown_frac = cooldown_frac
        # increase final validation ws, used for YaRN extension and short window size @classiclarryd
        self.ws_post_yarn_ext = ws_post_yarn_ext

        self.total_steps = self.scheduled_iterations + extension_iterations

        # Build stage boundaries (last is extension stage)
        ends = [0] + [round(c * scheduled_iterations) for c in accumulate(s.duration for s in stages[:-1])] + [self.total_steps]
        assert self.scheduled_iterations == ends[-2]
        self.boundaries = list(pairwise(ends))

        # Split embed at specified stage (ensure odd step for Adam)
        self.split_step = self.boundaries[split_embed_stage][0] | 1

        # Precompute MTP weights for all steps
        self.mtp_weights = []
        for step in range(self.total_steps + 1):
            stage, t = self.lookup(step)
            w = [a + (b - a) * t for a, b in zip(stage.mtp_weights_start, stage.mtp_weights_end)]
            self.mtp_weights.append(torch.tensor(w, device=device))

    def lookup(self, step: int) -> tuple[TrainingStage, float]:
        # Returns stage and % of the way through that stage
        for i, (start, end) in enumerate(self.boundaries):
            if step < end:
                t = (step - start) / (end - start)
                return self.stages[i], t
        return self.stages[-1], 1.0

    def get_lr(self, step: int) -> float:
        # learning rate schedule: tied to batch size schedule, with cooldown at the end
        stage, _ = self.lookup(step)
        lr = stage.lr_mul
        cd_start = int(self.scheduled_iterations * (1 - self.cooldown_frac))
        if step >= cd_start:
            t = min(1.0, (step - cd_start) / (self.scheduled_iterations - cd_start))
            lr = lr * (1 - t) + 0.1 * t
        return lr

# window_sizes are in units of `block_size` tokens (defined in TrainingManager)
TRAINING_STAGES = [
    TrainingStage(duration=1/3, batch_size=8 * 2048 * 8, window_sizes=(1, 3), lr_mul=1.0,
                  mtp_weights_start=[1.0, 0.5, 0.25], mtp_weights_end=[1.0, 0.5, 0.0]),
    TrainingStage(duration=1/3, batch_size=16 * 2048 * 8, window_sizes=(3, 7), lr_mul=1.52,  # (16/8)**0.6
                  mtp_weights_start=[1.0, 0.5], mtp_weights_end=[1.0, 0.0]),
    TrainingStage(duration=1/3, batch_size=24 * 2048 * 8, window_sizes=(5, 11), lr_mul=1.73,  # (24/8)**0.5
                  mtp_weights_start=[1.0], mtp_weights_end=[1.0]),
    # extension stage
    TrainingStage(batch_size=24 * 2048 * 8, window_sizes=(6, 13), lr_mul=1.0,  # lr_mul is not used
                  mtp_weights_start=[1.0], mtp_weights_end=[1.0]),
]

training_schedule = TrainingSchedule(TRAINING_STAGES, args.num_scheduled_iterations, args.num_extension_iterations, cooldown_frac=0.55)

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 = training_schedule.total_steps - 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.
        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
    """
    def __init__(self, model):
        self.model = model
        self.block_size = 128

        # - 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},
            "value_embed":    {"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
            "value_embed", "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 = training_schedule.split_step

        self.reset()

    def apply_final_ws_ext(self):
        self.ws_long = training_schedule.ws_post_yarn_ext

    def get_forward_args(self):
        return ForwardScheduleConfig(
            mtp_weights = self.mtp_weights,
            ws_short = self.ws_short * self.block_size,
            ws_long = self.ws_long * self.block_size
        )

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

    def get_transition_steps(self):
        return [start for start, _ in training_schedule.boundaries[1:]]

    def advance_schedule(self, step: int):
        stage, _ = training_schedule.lookup(step)
        self.ws_short, new_ws_long = stage.window_sizes
        if new_ws_long != self.ws_long:
            self.model.yarn.apply(self.ws_long * self.block_size, new_ws_long * self.block_size)
            self.model.yarn_paired_head.apply(self.ws_long * self.block_size, new_ws_long * self.block_size)

        new_batch_size = stage.batch_size
        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 = training_schedule.mtp_weights[step]

    def step_optimizers(self, step: int):
        step_lr = training_schedule.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()

        stage, _ = training_schedule.lookup(0)
        self.ws_short, self.ws_long = stage.window_sizes
        self.batch_size = stage.batch_size
        self.model.yarn.reset()
        self.model.yarn_paired_head.reset()

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

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

# 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, TRAINING_STAGES[0].batch_size, 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 = next(val_loader)
        model(inputs, targets, cum_seqlens, 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 = train_loader.send(send_args)
        (model(inputs, targets, cum_seqlens, training_manager.get_forward_args()) * grad_scale).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, TRAINING_STAGES[0].batch_size, 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 = training_schedule.total_steps
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 = next(val_loader)
                val_loss += model(inputs, targets, cum_seqlens, 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 = train_loader.send(training_manager.train_loader_send_args)
        (model(inputs, targets, cum_seqlens, training_manager.get_forward_args()) * grad_scale).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

@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.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

    # Hardcoded configs based on H100 autotuning
    if K == 768:
        BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K = 128, 128, 64
        num_stages, num_warps = 4, 4
    else:
        BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K = 64, 128, 128
        num_stages, num_warps = 4, 4

    grid = (batch_size * triton.cdiv(M, BLOCK_SIZE_M) * triton.cdiv(M, 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),
        BLOCK_SIZE_M=BLOCK_SIZE_M,
        BLOCK_SIZE_N=BLOCK_SIZE_N,
        BLOCK_SIZE_K=BLOCK_SIZE_K,
        GROUP_SIZE_M=8,
        LOWER_UPPER=1,
        num_stages=num_stages,
        num_warps=num_warps,
    )
    return out

@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

    # Hardcoded config based on H100 autotuning (M=768)
    BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K = 128, 128, 64
    num_stages, num_warps = 4, 4

    grid = (batch_size * triton.cdiv(M, BLOCK_SIZE_M) * triton.cdiv(M, 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,
        BLOCK_SIZE_M=BLOCK_SIZE_M,
        BLOCK_SIZE_N=BLOCK_SIZE_N,
        BLOCK_SIZE_K=BLOCK_SIZE_K,
        GROUP_SIZE_M=8,
        LOWER_UPPER=1,
        num_stages=num_stages,
        num_warps=num_warps,
    )
    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.12.7 (main, Jan 31 2026, 04:21:49) [GCC 13.2.0]
Running PyTorch 2.10.0.dev20251210+cu126 compiled for CUDA 12.6
Running Triton version 3.6.0
Sun Feb  1 06:10:04 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:63:00.0 Off |                    0 |
| N/A   33C    P0            117W /  700W |    1519MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   1  NVIDIA H100 80GB HBM3          On  |   00000000:6B:00.0 Off |                    0 |
| N/A   37C    P0            124W /  700W |    1519MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   2  NVIDIA H100 80GB HBM3          On  |   00000000:71:00.0 Off |                    0 |
| N/A   39C    P0            124W /  700W |    1519MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   3  NVIDIA H100 80GB HBM3          On  |   00000000:79:00.0 Off |                    0 |
| N/A   34C    P0            126W /  700W |    1519MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   4  NVIDIA H100 80GB HBM3          On  |   00000000:7F:00.0 Off |                    0 |
| N/A   33C    P0            119W /  700W |    1519MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   5  NVIDIA H100 80GB HBM3          On  |   00000000:87:00.0 Off |                    0 |
| N/A   39C    P0            120W /  700W |    1519MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   6  NVIDIA H100 80GB HBM3          On  |   00000000:8D:00.0 Off |                    0 |
| N/A   37C    P0            124W /  700W |    1519MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   7  NVIDIA H100 80GB HBM3          On  |   00000000:95:00.0 Off |                    0 |
| N/A   34C    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           18107      C   /usr/local/bin/python                  1510MiB |
|    1   N/A  N/A           18108      C   /usr/local/bin/python                  1510MiB |
|    2   N/A  N/A           18109      C   /usr/local/bin/python                  1510MiB |
|    3   N/A  N/A           18110      C   /usr/local/bin/python                  1510MiB |
|    4   N/A  N/A           18111      C   /usr/local/bin/python                  1510MiB |
|    5   N/A  N/A           18112      C   /usr/local/bin/python                  1510MiB |
|    6   N/A  N/A           18113      C   /usr/local/bin/python                  1510MiB |
|    7   N/A  N/A           18114      C   /usr/local/bin/python                  1510MiB |
+-----------------------------------------------------------------------------------------+

====================================================================================================
Compiling model and warming up kernels (~7 minutes on first execution)
Sampling steps [0, 1, 2, 504, 505, 506, 1009, 1010, 1011, 1514, 1515, 1516] for warmup
Resetting Model
step:0/1555 val_loss:10.8306 train_time:0ms step_avg:0.03ms
step:1/1555 train_time:74ms step_avg:74.30ms
step:2/1555 train_time:95ms step_avg:47.73ms
step:3/1555 train_time:119ms step_avg:39.57ms
step:4/1555 train_time:156ms step_avg:39.10ms
step:5/1555 train_time:187ms step_avg:37.37ms
step:6/1555 train_time:224ms step_avg:37.38ms
step:7/1555 train_time:255ms step_avg:36.43ms
step:8/1555 train_time:293ms step_avg:36.61ms
step:9/1555 train_time:324ms step_avg:36.01ms
step:10/1555 train_time:361ms step_avg:36.13ms
step:11/1555 train_time:392ms step_avg:35.67ms
step:12/1555 train_time:430ms step_avg:35.85ms
step:13/1555 train_time:461ms step_avg:35.45ms
step:14/1555 train_time:498ms step_avg:35.59ms
step:15/1555 train_time:529ms step_avg:35.29ms
step:16/1555 train_time:568ms step_avg:35.47ms
step:17/1555 train_time:599ms step_avg:35.24ms
step:18/1555 train_time:637ms step_avg:35.36ms
step:19/1555 train_time:667ms step_avg:35.13ms
step:20/1555 train_time:705ms step_avg:35.23ms
step:21/1555 train_time:736ms step_avg:35.05ms
step:22/1555 train_time:774ms step_avg:35.17ms
step:23/1555 train_time:805ms step_avg:34.99ms
step:24/1555 train_time:842ms step_avg:35.08ms
step:25/1555 train_time:873ms step_avg:34.93ms
step:26/1555 train_time:911ms step_avg:35.04ms
step:27/1555 train_time:942ms step_avg:34.89ms
step:28/1555 train_time:980ms step_avg:34.99ms
step:29/1555 train_time:1011ms step_avg:34.85ms
step:30/1555 train_time:1048ms step_avg:34.93ms
step:31/1555 train_time:1079ms step_avg:34.82ms
step:32/1555 train_time:1117ms step_avg:34.90ms
step:33/1555 train_time:1148ms step_avg:34.80ms
step:34/1555 train_time:1186ms step_avg:34.88ms
step:35/1555 train_time:1217ms step_avg:34.78ms
step:36/1555 train_time:1255ms step_avg:34.86ms
step:37/1555 train_time:1286ms step_avg:34.75ms
step:38/1555 train_time:1323ms step_avg:34.82ms
step:39/1555 train_time:1355ms step_avg:34.73ms
step:40/1555 train_time:1392ms step_avg:34.80ms
step:41/1555 train_time:1423ms step_avg:34.72ms
step:42/1555 train_time:1461ms step_avg:34.78ms
step:43/1555 train_time:1492ms step_avg:34.70ms
step:44/1555 train_time:1530ms step_avg:34.77ms
step:45/1555 train_time:1562ms step_avg:34.70ms
step:46/1555 train_time:1599ms step_avg:34.76ms
step:47/1555 train_time:1630ms step_avg:34.68ms
step:48/1555 train_time:1669ms step_avg:34.76ms
step:49/1555 train_time:1699ms step_avg:34.68ms
step:50/1555 train_time:1737ms step_avg:34.75ms
step:51/1555 train_time:1769ms step_avg:34.68ms
step:52/1555 train_time:1807ms step_avg:34.74ms
step:53/1555 train_time:1838ms step_avg:34.67ms
step:54/1555 train_time:1875ms step_avg:34.73ms
step:55/1555 train_time:1906ms step_avg:34.66ms
step:56/1555 train_time:1944ms step_avg:34.71ms
step:57/1555 train_time:1975ms step_avg:34.65ms
step:58/1555 train_time:2013ms step_avg:34.71ms
step:59/1555 train_time:2045ms step_avg:34.65ms
step:60/1555 train_time:2082ms step_avg:34.70ms
step:61/1555 train_time:2113ms step_avg:34.64ms
step:62/1555 train_time:2151ms step_avg:34.70ms
step:63/1555 train_time:2182ms step_avg:34.64ms
step:64/1555 train_time:2220ms step_avg:34.68ms
step:65/1555 train_time:2251ms step_avg:34.63ms
step:66/1555 train_time:2289ms step_avg:34.68ms
step:67/1555 train_time:2320ms step_avg:34.63ms
step:68/1555 train_time:2357ms step_avg:34.67ms
step:69/1555 train_time:2389ms step_avg:34.62ms
step:70/1555 train_time:2426ms step_avg:34.66ms
step:71/1555 train_time:2457ms step_avg:34.61ms
step:72/1555 train_time:2495ms step_avg:34.65ms
step:73/1555 train_time:2526ms step_avg:34.60ms
step:74/1555 train_time:2563ms step_avg:34.63ms
step:75/1555 train_time:2594ms step_avg:34.58ms
step:76/1555 train_time:2631ms step_avg:34.62ms
step:77/1555 train_time:2662ms step_avg:34.58ms
step:78/1555 train_time:2700ms step_avg:34.61ms
step:79/1555 train_time:2731ms step_avg:34.57ms
step:80/1555 train_time:2769ms step_avg:34.61ms
step:81/1555 train_time:2800ms step_avg:34.57ms
step:82/1555 train_time:2838ms step_avg:34.60ms
step:83/1555 train_time:2869ms step_avg:34.56ms
step:84/1555 train_time:2906ms step_avg:34.60ms
step:85/1555 train_time:2938ms step_avg:34.56ms
step:86/1555 train_time:2976ms step_avg:34.60ms
step:87/1555 train_time:3007ms step_avg:34.56ms
step:88/1555 train_time:3044ms step_avg:34.59ms
step:89/1555 train_time:3076ms step_avg:34.56ms
step:90/1555 train_time:3113ms step_avg:34.59ms
step:91/1555 train_time:3144ms step_avg:34.55ms
step:92/1555 train_time:3181ms step_avg:34.58ms
step:93/1555 train_time:3213ms step_avg:34.55ms
step:94/1555 train_time:3250ms step_avg:34.58ms
step:95/1555 train_time:3281ms step_avg:34.54ms
step:96/1555 train_time:3319ms step_avg:34.57ms
step:97/1555 train_time:3350ms step_avg:34.53ms
step:98/1555 train_time:3387ms step_avg:34.57ms
step:99/1555 train_time:3419ms step_avg:34.54ms
step:100/1555 train_time:3457ms step_avg:34.57ms
step:101/1555 train_time:3488ms step_avg:34.53ms
step:102/1555 train_time:3525ms step_avg:34.56ms
step:103/1555 train_time:3557ms step_avg:34.53ms
step:104/1555 train_time:3594ms step_avg:34.56ms
step:105/1555 train_time:3626ms step_avg:34.53ms
step:106/1555 train_time:3663ms step_avg:34.56ms
step:107/1555 train_time:3694ms step_avg:34.53ms
step:108/1555 train_time:3732ms step_avg:34.56ms
step:109/1555 train_time:3763ms step_avg:34.52ms
step:110/1555 train_time:3800ms step_avg:34.55ms
step:111/1555 train_time:3831ms step_avg:34.52ms
step:112/1555 train_time:3869ms step_avg:34.54ms
step:113/1555 train_time:3900ms step_avg:34.51ms
step:114/1555 train_time:3937ms step_avg:34.54ms
step:115/1555 train_time:3969ms step_avg:34.51ms
step:116/1555 train_time:4006ms step_avg:34.53ms
step:117/1555 train_time:4037ms step_avg:34.50ms
step:118/1555 train_time:4075ms step_avg:34.53ms
step:119/1555 train_time:4106ms step_avg:34.50ms
step:120/1555 train_time:4143ms step_avg:34.53ms
step:121/1555 train_time:4175ms step_avg:34.50ms
step:122/1555 train_time:4213ms step_avg:34.53ms
step:123/1555 train_time:4244ms step_avg:34.50ms
step:124/1555 train_time:4281ms step_avg:34.52ms
step:125/1555 train_time:4313ms step_avg:34.50ms
step:126/1555 train_time:4350ms step_avg:34.53ms
step:127/1555 train_time:4381ms step_avg:34.50ms
step:128/1555 train_time:4418ms step_avg:34.52ms
step:129/1555 train_time:4450ms step_avg:34.49ms
step:130/1555 train_time:4488ms step_avg:34.53ms
step:131/1555 train_time:4520ms step_avg:34.50ms
step:132/1555 train_time:4557ms step_avg:34.52ms
step:133/1555 train_time:4588ms step_avg:34.50ms
step:134/1555 train_time:4626ms step_avg:34.52ms
step:135/1555 train_time:4657ms step_avg:34.50ms
step:136/1555 train_time:4695ms step_avg:34.52ms
step:137/1555 train_time:4726ms step_avg:34.49ms
step:138/1555 train_time:4763ms step_avg:34.51ms
step:139/1555 train_time:4794ms step_avg:34.49ms
step:140/1555 train_time:4832ms step_avg:34.51ms
step:141/1555 train_time:4862ms step_avg:34.49ms
step:142/1555 train_time:4900ms step_avg:34.50ms
step:143/1555 train_time:4931ms step_avg:34.48ms
step:144/1555 train_time:4968ms step_avg:34.50ms
step:145/1555 train_time:4999ms step_avg:34.48ms
step:146/1555 train_time:5037ms step_avg:34.50ms
step:147/1555 train_time:5068ms step_avg:34.48ms
step:148/1555 train_time:5105ms step_avg:34.49ms
step:149/1555 train_time:5136ms step_avg:34.47ms
step:150/1555 train_time:5174ms step_avg:34.49ms
step:151/1555 train_time:5205ms step_avg:34.47ms
step:152/1555 train_time:5242ms step_avg:34.49ms
step:153/1555 train_time:5274ms step_avg:34.47ms
step:154/1555 train_time:5311ms step_avg:34.49ms
step:155/1555 train_time:5342ms step_avg:34.47ms
step:156/1555 train_time:5380ms step_avg:34.49ms
step:157/1555 train_time:5411ms step_avg:34.47ms
step:158/1555 train_time:5450ms step_avg:34.49ms
step:159/1555 train_time:5481ms step_avg:34.47ms
step:160/1555 train_time:5518ms step_avg:34.49ms
step:161/1555 train_time:5550ms step_avg:34.47ms
step:162/1555 train_time:5587ms step_avg:34.49ms
step:163/1555 train_time:5619ms step_avg:34.47ms
step:164/1555 train_time:5656ms step_avg:34.49ms
step:165/1555 train_time:5688ms step_avg:34.47ms
step:166/1555 train_time:5725ms step_avg:34.49ms
step:167/1555 train_time:5756ms step_avg:34.47ms
step:168/1555 train_time:5794ms step_avg:34.49ms
step:169/1555 train_time:5826ms step_avg:34.47ms
step:170/1555 train_time:5863ms step_avg:34.49ms
step:171/1555 train_time:5894ms step_avg:34.47ms
step:172/1555 train_time:5931ms step_avg:34.48ms
step:173/1555 train_time:5962ms step_avg:34.46ms
step:174/1555 train_time:6000ms step_avg:34.48ms
step:175/1555 train_time:6031ms step_avg:34.46ms
step:176/1555 train_time:6069ms step_avg:34.48ms
step:177/1555 train_time:6100ms step_avg:34.46ms
step:178/1555 train_time:6137ms step_avg:34.48ms
step:179/1555 train_time:6169ms step_avg:34.46ms
step:180/1555 train_time:6206ms step_avg:34.48ms
step:181/1555 train_time:6237ms step_avg:34.46ms
step:182/1555 train_time:6275ms step_avg:34.48ms
step:183/1555 train_time:6306ms step_avg:34.46ms
step:184/1555 train_time:6343ms step_avg:34.47ms
step:185/1555 train_time:6374ms step_avg:34.45ms
step:186/1555 train_time:6412ms step_avg:34.47ms
step:187/1555 train_time:6443ms step_avg:34.45ms
step:188/1555 train_time:6480ms step_avg:34.47ms
step:189/1555 train_time:6511ms step_avg:34.45ms
step:190/1555 train_time:6548ms step_avg:34.47ms
step:191/1555 train_time:6579ms step_avg:34.45ms
step:192/1555 train_time:6617ms step_avg:34.46ms
step:193/1555 train_time:6648ms step_avg:34.45ms
step:194/1555 train_time:6686ms step_avg:34.46ms
step:195/1555 train_time:6717ms step_avg:34.44ms
step:196/1555 train_time:6755ms step_avg:34.46ms
step:197/1555 train_time:6786ms step_avg:34.44ms
step:198/1555 train_time:6823ms step_avg:34.46ms
step:199/1555 train_time:6854ms step_avg:34.44ms
step:200/1555 train_time:6892ms step_avg:34.46ms
step:201/1555 train_time:6923ms step_avg:34.44ms
step:202/1555 train_time:6960ms step_avg:34.45ms
step:203/1555 train_time:6991ms step_avg:34.44ms
step:204/1555 train_time:7028ms step_avg:34.45ms
step:205/1555 train_time:7059ms step_avg:34.44ms
step:206/1555 train_time:7097ms step_avg:34.45ms
step:207/1555 train_time:7128ms step_avg:34.43ms
step:208/1555 train_time:7165ms step_avg:34.45ms
step:209/1555 train_time:7196ms step_avg:34.43ms
step:210/1555 train_time:7234ms step_avg:34.45ms
step:211/1555 train_time:7265ms step_avg:34.43ms
step:212/1555 train_time:7303ms step_avg:34.45ms
step:213/1555 train_time:7334ms step_avg:34.43ms
step:214/1555 train_time:7372ms step_avg:34.45ms
step:215/1555 train_time:7404ms step_avg:34.44ms
step:216/1555 train_time:7441ms step_avg:34.45ms
step:217/1555 train_time:7472ms step_avg:34.43ms
step:218/1555 train_time:7509ms step_avg:34.45ms
step:219/1555 train_time:7540ms step_avg:34.43ms
step:220/1555 train_time:7577ms step_avg:34.44ms
step:221/1555 train_time:7609ms step_avg:34.43ms
step:222/1555 train_time:7647ms step_avg:34.44ms
step:223/1555 train_time:7678ms step_avg:34.43ms
step:224/1555 train_time:7715ms step_avg:34.44ms
step:225/1555 train_time:7746ms step_avg:34.43ms
step:226/1555 train_time:7783ms step_avg:34.44ms
step:227/1555 train_time:7814ms step_avg:34.42ms
step:228/1555 train_time:7852ms step_avg:34.44ms
step:229/1555 train_time:7883ms step_avg:34.43ms
step:230/1555 train_time:7921ms step_avg:34.44ms
step:231/1555 train_time:7952ms step_avg:34.42ms
step:232/1555 train_time:7990ms step_avg:34.44ms
step:233/1555 train_time:8020ms step_avg:34.42ms
step:234/1555 train_time:8058ms step_avg:34.43ms
step:235/1555 train_time:8088ms step_avg:34.42ms
step:236/1555 train_time:8126ms step_avg:34.43ms
step:237/1555 train_time:8156ms step_avg:34.42ms
step:238/1555 train_time:8194ms step_avg:34.43ms
step:239/1555 train_time:8225ms step_avg:34.41ms
step:240/1555 train_time:8262ms step_avg:34.42ms
step:241/1555 train_time:8293ms step_avg:34.41ms
step:242/1555 train_time:8330ms step_avg:34.42ms
step:243/1555 train_time:8361ms step_avg:34.41ms
step:244/1555 train_time:8399ms step_avg:34.42ms
step:245/1555 train_time:8430ms step_avg:34.41ms
step:246/1555 train_time:8467ms step_avg:34.42ms
step:247/1555 train_time:8498ms step_avg:34.41ms
step:248/1555 train_time:8536ms step_avg:34.42ms
step:249/1555 train_time:8568ms step_avg:34.41ms
step:250/1555 train_time:8605ms step_avg:34.42ms
step:250/1555 val_loss:4.5542 train_time:8656ms step_avg:34.62ms
step:251/1555 train_time:8673ms step_avg:34.55ms
step:252/1555 train_time:8692ms step_avg:34.49ms
step:253/1555 train_time:8709ms step_avg:34.42ms
step:254/1555 train_time:8745ms step_avg:34.43ms
step:255/1555 train_time:8778ms step_avg:34.42ms
step:256/1555 train_time:8816ms step_avg:34.44ms
step:257/1555 train_time:8848ms step_avg:34.43ms
step:258/1555 train_time:8886ms step_avg:34.44ms
step:259/1555 train_time:8917ms step_avg:34.43ms
step:260/1555 train_time:8954ms step_avg:34.44ms
step:261/1555 train_time:8986ms step_avg:34.43ms
step:262/1555 train_time:9024ms step_avg:34.44ms
step:263/1555 train_time:9055ms step_avg:34.43ms
step:264/1555 train_time:9092ms step_avg:34.44ms
step:265/1555 train_time:9124ms step_avg:34.43ms
step:266/1555 train_time:9161ms step_avg:34.44ms
step:267/1555 train_time:9192ms step_avg:34.43ms
step:268/1555 train_time:9229ms step_avg:34.44ms
step:269/1555 train_time:9260ms step_avg:34.42ms
step:270/1555 train_time:9297ms step_avg:34.43ms
step:271/1555 train_time:9328ms step_avg:34.42ms
step:272/1555 train_time:9365ms step_avg:34.43ms
step:273/1555 train_time:9396ms step_avg:34.42ms
step:274/1555 train_time:9433ms step_avg:34.43ms
step:275/1555 train_time:9464ms step_avg:34.41ms
step:276/1555 train_time:9502ms step_avg:34.43ms
step:277/1555 train_time:9533ms step_avg:34.41ms
step:278/1555 train_time:9570ms step_avg:34.42ms
step:279/1555 train_time:9601ms step_avg:34.41ms
step:280/1555 train_time:9638ms step_avg:34.42ms
step:281/1555 train_time:9669ms step_avg:34.41ms
step:282/1555 train_time:9707ms step_avg:34.42ms
step:283/1555 train_time:9738ms step_avg:34.41ms
step:284/1555 train_time:9775ms step_avg:34.42ms
step:285/1555 train_time:9806ms step_avg:34.41ms
step:286/1555 train_time:9844ms step_avg:34.42ms
step:287/1555 train_time:9876ms step_avg:34.41ms
step:288/1555 train_time:9913ms step_avg:34.42ms
step:289/1555 train_time:9944ms step_avg:34.41ms
step:290/1555 train_time:9982ms step_avg:34.42ms
step:291/1555 train_time:10013ms step_avg:34.41ms
step:292/1555 train_time:10050ms step_avg:34.42ms
step:293/1555 train_time:10081ms step_avg:34.41ms
step:294/1555 train_time:10119ms step_avg:34.42ms
step:295/1555 train_time:10150ms step_avg:34.41ms
step:296/1555 train_time:10187ms step_avg:34.42ms
step:297/1555 train_time:10218ms step_avg:34.40ms
step:298/1555 train_time:10255ms step_avg:34.41ms
step:299/1555 train_time:10286ms step_avg:34.40ms
step:300/1555 train_time:10324ms step_avg:34.41ms
step:301/1555 train_time:10355ms step_avg:34.40ms
step:302/1555 train_time:10392ms step_avg:34.41ms
step:303/1555 train_time:10423ms step_avg:34.40ms
step:304/1555 train_time:10460ms step_avg:34.41ms
step:305/1555 train_time:10491ms step_avg:34.40ms
step:306/1555 train_time:10528ms step_avg:34.41ms
step:307/1555 train_time:10559ms step_avg:34.39ms
step:308/1555 train_time:10596ms step_avg:34.40ms
step:309/1555 train_time:10627ms step_avg:34.39ms
step:310/1555 train_time:10664ms step_avg:34.40ms
step:311/1555 train_time:10695ms step_avg:34.39ms
step:312/1555 train_time:10732ms step_avg:34.40ms
step:313/1555 train_time:10763ms step_avg:34.39ms
step:314/1555 train_time:10800ms step_avg:34.40ms
step:315/1555 train_time:10831ms step_avg:34.38ms
step:316/1555 train_time:10869ms step_avg:34.39ms
step:317/1555 train_time:10900ms step_avg:34.38ms
step:318/1555 train_time:10937ms step_avg:34.39ms
step:319/1555 train_time:10969ms step_avg:34.39ms
step:320/1555 train_time:11006ms step_avg:34.39ms
step:321/1555 train_time:11037ms step_avg:34.38ms
step:322/1555 train_time:11074ms step_avg:34.39ms
step:323/1555 train_time:11105ms step_avg:34.38ms
step:324/1555 train_time:11143ms step_avg:34.39ms
step:325/1555 train_time:11174ms step_avg:34.38ms
step:326/1555 train_time:11211ms step_avg:34.39ms
step:327/1555 train_time:11243ms step_avg:34.38ms
step:328/1555 train_time:11280ms step_avg:34.39ms
step:329/1555 train_time:11311ms step_avg:34.38ms
step:330/1555 train_time:11349ms step_avg:34.39ms
step:331/1555 train_time:11380ms step_avg:34.38ms
step:332/1555 train_time:11417ms step_avg:34.39ms
step:333/1555 train_time:11448ms step_avg:34.38ms
step:334/1555 train_time:11486ms step_avg:34.39ms
step:335/1555 train_time:11517ms step_avg:34.38ms
step:336/1555 train_time:11554ms step_avg:34.39ms
step:337/1555 train_time:11586ms step_avg:34.38ms
step:338/1555 train_time:11624ms step_avg:34.39ms
step:339/1555 train_time:11655ms step_avg:34.38ms
step:340/1555 train_time:11692ms step_avg:34.39ms
step:341/1555 train_time:11723ms step_avg:34.38ms
step:342/1555 train_time:11760ms step_avg:34.39ms
step:343/1555 train_time:11791ms step_avg:34.38ms
step:344/1555 train_time:11828ms step_avg:34.38ms
step:345/1555 train_time:11859ms step_avg:34.37ms
step:346/1555 train_time:11896ms step_avg:34.38ms
step:347/1555 train_time:11927ms step_avg:34.37ms
step:348/1555 train_time:11965ms step_avg:34.38ms
step:349/1555 train_time:11995ms step_avg:34.37ms
step:350/1555 train_time:12033ms step_avg:34.38ms
step:351/1555 train_time:12063ms step_avg:34.37ms
step:352/1555 train_time:12101ms step_avg:34.38ms
step:353/1555 train_time:12132ms step_avg:34.37ms
step:354/1555 train_time:12169ms step_avg:34.38ms
step:355/1555 train_time:12200ms step_avg:34.37ms
step:356/1555 train_time:12237ms step_avg:34.37ms
step:357/1555 train_time:12269ms step_avg:34.37ms
step:358/1555 train_time:12306ms step_avg:34.38ms
step:359/1555 train_time:12338ms step_avg:34.37ms
step:360/1555 train_time:12374ms step_avg:34.37ms
step:361/1555 train_time:12405ms step_avg:34.36ms
step:362/1555 train_time:12443ms step_avg:34.37ms
step:363/1555 train_time:12474ms step_avg:34.36ms
step:364/1555 train_time:12511ms step_avg:34.37ms
step:365/1555 train_time:12542ms step_avg:34.36ms
step:366/1555 train_time:12580ms step_avg:34.37ms
step:367/1555 train_time:12611ms step_avg:34.36ms
step:368/1555 train_time:12649ms step_avg:34.37ms
step:369/1555 train_time:12680ms step_avg:34.36ms
step:370/1555 train_time:12717ms step_avg:34.37ms
step:371/1555 train_time:12748ms step_avg:34.36ms
step:372/1555 train_time:12786ms step_avg:34.37ms
step:373/1555 train_time:12816ms step_avg:34.36ms
step:374/1555 train_time:12853ms step_avg:34.37ms
step:375/1555 train_time:12885ms step_avg:34.36ms
step:376/1555 train_time:12922ms step_avg:34.37ms
step:377/1555 train_time:12953ms step_avg:34.36ms
step:378/1555 train_time:12990ms step_avg:34.37ms
step:379/1555 train_time:13021ms step_avg:34.36ms
step:380/1555 train_time:13059ms step_avg:34.37ms
step:381/1555 train_time:13090ms step_avg:34.36ms
step:382/1555 train_time:13127ms step_avg:34.36ms
step:383/1555 train_time:13158ms step_avg:34.35ms
step:384/1555 train_time:13195ms step_avg:34.36ms
step:385/1555 train_time:13226ms step_avg:34.35ms
step:386/1555 train_time:13263ms step_avg:34.36ms
step:387/1555 train_time:13294ms step_avg:34.35ms
step:388/1555 train_time:13331ms step_avg:34.36ms
step:389/1555 train_time:13363ms step_avg:34.35ms
step:390/1555 train_time:13400ms step_avg:34.36ms
step:391/1555 train_time:13431ms step_avg:34.35ms
step:392/1555 train_time:13468ms step_avg:34.36ms
step:393/1555 train_time:13499ms step_avg:34.35ms
step:394/1555 train_time:13536ms step_avg:34.36ms
step:395/1555 train_time:13568ms step_avg:34.35ms
step:396/1555 train_time:13606ms step_avg:34.36ms
step:397/1555 train_time:13636ms step_avg:34.35ms
step:398/1555 train_time:13674ms step_avg:34.36ms
step:399/1555 train_time:13705ms step_avg:34.35ms
step:400/1555 train_time:13743ms step_avg:34.36ms
step:401/1555 train_time:13774ms step_avg:34.35ms
step:402/1555 train_time:13812ms step_avg:34.36ms
step:403/1555 train_time:13843ms step_avg:34.35ms
step:404/1555 train_time:13881ms step_avg:34.36ms
step:405/1555 train_time:13912ms step_avg:34.35ms
step:406/1555 train_time:13949ms step_avg:34.36ms
step:407/1555 train_time:13980ms step_avg:34.35ms
step:408/1555 train_time:14017ms step_avg:34.36ms
step:409/1555 train_time:14049ms step_avg:34.35ms
step:410/1555 train_time:14086ms step_avg:34.36ms
step:411/1555 train_time:14117ms step_avg:34.35ms
step:412/1555 train_time:14154ms step_avg:34.36ms
step:413/1555 train_time:14185ms step_avg:34.35ms
step:414/1555 train_time:14222ms step_avg:34.35ms
step:415/1555 train_time:14253ms step_avg:34.35ms
step:416/1555 train_time:14290ms step_avg:34.35ms
step:417/1555 train_time:14321ms step_avg:34.34ms
step:418/1555 train_time:14358ms step_avg:34.35ms
step:419/1555 train_time:14389ms step_avg:34.34ms
step:420/1555 train_time:14427ms step_avg:34.35ms
step:421/1555 train_time:14458ms step_avg:34.34ms
step:422/1555 train_time:14495ms step_avg:34.35ms
step:423/1555 train_time:14526ms step_avg:34.34ms
step:424/1555 train_time:14563ms step_avg:34.35ms
step:425/1555 train_time:14594ms step_avg:34.34ms
step:426/1555 train_time:14632ms step_avg:34.35ms
step:427/1555 train_time:14662ms step_avg:34.34ms
step:428/1555 train_time:14700ms step_avg:34.34ms
step:429/1555 train_time:14731ms step_avg:34.34ms
step:430/1555 train_time:14768ms step_avg:34.34ms
step:431/1555 train_time:14799ms step_avg:34.34ms
step:432/1555 train_time:14836ms step_avg:34.34ms
step:433/1555 train_time:14867ms step_avg:34.34ms
step:434/1555 train_time:14905ms step_avg:34.34ms
step:435/1555 train_time:14936ms step_avg:34.34ms
step:436/1555 train_time:14973ms step_avg:34.34ms
step:437/1555 train_time:15004ms step_avg:34.33ms
step:438/1555 train_time:15041ms step_avg:34.34ms
step:439/1555 train_time:15073ms step_avg:34.33ms
step:440/1555 train_time:15110ms step_avg:34.34ms
step:441/1555 train_time:15141ms step_avg:34.33ms
step:442/1555 train_time:15179ms step_avg:34.34ms
step:443/1555 train_time:15209ms step_avg:34.33ms
step:444/1555 train_time:15247ms step_avg:34.34ms
step:445/1555 train_time:15278ms step_avg:34.33ms
step:446/1555 train_time:15315ms step_avg:34.34ms
step:447/1555 train_time:15346ms step_avg:34.33ms
step:448/1555 train_time:15383ms step_avg:34.34ms
step:449/1555 train_time:15414ms step_avg:34.33ms
step:450/1555 train_time:15452ms step_avg:34.34ms
step:451/1555 train_time:15483ms step_avg:34.33ms
step:452/1555 train_time:15521ms step_avg:34.34ms
step:453/1555 train_time:15552ms step_avg:34.33ms
step:454/1555 train_time:15589ms step_avg:34.34ms
step:455/1555 train_time:15620ms step_avg:34.33ms
step:456/1555 train_time:15657ms step_avg:34.34ms
step:457/1555 train_time:15688ms step_avg:34.33ms
step:458/1555 train_time:15725ms step_avg:34.33ms
step:459/1555 train_time:15756ms step_avg:34.33ms
step:460/1555 train_time:15793ms step_avg:34.33ms
step:461/1555 train_time:15824ms step_avg:34.33ms
step:462/1555 train_time:15862ms step_avg:34.33ms
step:463/1555 train_time:15893ms step_avg:34.33ms
step:464/1555 train_time:15930ms step_avg:34.33ms
step:465/1555 train_time:15961ms step_avg:34.32ms
step:466/1555 train_time:15998ms step_avg:34.33ms
step:467/1555 train_time:16029ms step_avg:34.32ms
step:468/1555 train_time:16066ms step_avg:34.33ms
step:469/1555 train_time:16097ms step_avg:34.32ms
step:470/1555 train_time:16134ms step_avg:34.33ms
step:471/1555 train_time:16165ms step_avg:34.32ms
step:472/1555 train_time:16202ms step_avg:34.33ms
step:473/1555 train_time:16233ms step_avg:34.32ms
step:474/1555 train_time:16270ms step_avg:34.33ms
step:475/1555 train_time:16302ms step_avg:34.32ms
step:476/1555 train_time:16339ms step_avg:34.32ms
step:477/1555 train_time:16370ms step_avg:34.32ms
step:478/1555 train_time:16407ms step_avg:34.32ms
step:479/1555 train_time:16438ms step_avg:34.32ms
step:480/1555 train_time:16476ms step_avg:34.32ms
step:481/1555 train_time:16506ms step_avg:34.32ms
step:482/1555 train_time:16544ms step_avg:34.32ms
step:483/1555 train_time:16575ms step_avg:34.32ms
step:484/1555 train_time:16612ms step_avg:34.32ms
step:485/1555 train_time:16643ms step_avg:34.32ms
step:486/1555 train_time:16681ms step_avg:34.32ms
step:487/1555 train_time:16712ms step_avg:34.32ms
step:488/1555 train_time:16749ms step_avg:34.32ms
step:489/1555 train_time:16780ms step_avg:34.32ms
step:490/1555 train_time:16817ms step_avg:34.32ms
step:491/1555 train_time:16848ms step_avg:34.31ms
step:492/1555 train_time:16886ms step_avg:34.32ms
step:493/1555 train_time:16917ms step_avg:34.31ms
step:494/1555 train_time:16954ms step_avg:34.32ms
step:495/1555 train_time:16985ms step_avg:34.31ms
step:496/1555 train_time:17023ms step_avg:34.32ms
step:497/1555 train_time:17055ms step_avg:34.31ms
step:498/1555 train_time:17092ms step_avg:34.32ms
step:499/1555 train_time:17123ms step_avg:34.31ms
step:500/1555 train_time:17160ms step_avg:34.32ms
step:500/1555 val_loss:4.2227 train_time:17210ms step_avg:34.42ms
step:501/1555 train_time:17227ms step_avg:34.39ms
step:502/1555 train_time:17246ms step_avg:34.36ms
step:503/1555 train_time:17263ms step_avg:34.32ms
step:504/1555 train_time:17299ms step_avg:34.32ms
step:505/1555 train_time:17331ms step_avg:34.32ms
step:506/1555 train_time:17372ms step_avg:34.33ms
step:507/1555 train_time:17427ms step_avg:34.37ms
step:508/1555 train_time:17491ms step_avg:34.43ms
step:509/1555 train_time:17549ms step_avg:34.48ms
step:510/1555 train_time:17613ms step_avg:34.53ms
step:511/1555 train_time:17671ms step_avg:34.58ms
step:512/1555 train_time:17734ms step_avg:34.64ms
step:513/1555 train_time:17790ms step_avg:34.68ms
step:514/1555 train_time:17854ms step_avg:34.74ms
step:515/1555 train_time:17910ms step_avg:34.78ms
step:516/1555 train_time:17973ms step_avg:34.83ms
step:517/1555 train_time:18030ms step_avg:34.87ms
step:518/1555 train_time:18093ms step_avg:34.93ms
step:519/1555 train_time:18152ms step_avg:34.98ms
step:520/1555 train_time:18217ms step_avg:35.03ms
step:521/1555 train_time:18276ms step_avg:35.08ms
step:522/1555 train_time:18341ms step_avg:35.14ms
step:523/1555 train_time:18399ms step_avg:35.18ms
step:524/1555 train_time:18464ms step_avg:35.24ms
step:525/1555 train_time:18522ms step_avg:35.28ms
step:526/1555 train_time:18587ms step_avg:35.34ms
step:527/1555 train_time:18644ms step_avg:35.38ms
step:528/1555 train_time:18708ms step_avg:35.43ms
step:529/1555 train_time:18765ms step_avg:35.47ms
step:530/1555 train_time:18829ms step_avg:35.53ms
step:531/1555 train_time:18886ms step_avg:35.57ms
step:532/1555 train_time:18950ms step_avg:35.62ms
step:533/1555 train_time:19007ms step_avg:35.66ms
step:534/1555 train_time:19071ms step_avg:35.71ms
step:535/1555 train_time:19130ms step_avg:35.76ms
step:536/1555 train_time:19195ms step_avg:35.81ms
step:537/1555 train_time:19254ms step_avg:35.86ms
step:538/1555 train_time:19318ms step_avg:35.91ms
step:539/1555 train_time:19376ms step_avg:35.95ms
step:540/1555 train_time:19439ms step_avg:36.00ms
step:541/1555 train_time:19497ms step_avg:36.04ms
step:542/1555 train_time:19561ms step_avg:36.09ms
step:543/1555 train_time:19619ms step_avg:36.13ms
step:544/1555 train_time:19683ms step_avg:36.18ms
step:545/1555 train_time:19741ms step_avg:36.22ms
step:546/1555 train_time:19805ms step_avg:36.27ms
step:547/1555 train_time:19862ms step_avg:36.31ms
step:548/1555 train_time:19926ms step_avg:36.36ms
step:549/1555 train_time:19984ms step_avg:36.40ms
step:550/1555 train_time:20048ms step_avg:36.45ms
step:551/1555 train_time:20106ms step_avg:36.49ms
step:552/1555 train_time:20171ms step_avg:36.54ms
step:553/1555 train_time:20230ms step_avg:36.58ms
step:554/1555 train_time:20294ms step_avg:36.63ms
step:555/1555 train_time:20352ms step_avg:36.67ms
step:556/1555 train_time:20416ms step_avg:36.72ms
step:557/1555 train_time:20474ms step_avg:36.76ms
step:558/1555 train_time:20538ms step_avg:36.81ms
step:559/1555 train_time:20595ms step_avg:36.84ms
step:560/1555 train_time:20660ms step_avg:36.89ms
step:561/1555 train_time:20716ms step_avg:36.93ms
step:562/1555 train_time:20781ms step_avg:36.98ms
step:563/1555 train_time:20838ms step_avg:37.01ms
step:564/1555 train_time:20902ms step_avg:37.06ms
step:565/1555 train_time:20960ms step_avg:37.10ms
step:566/1555 train_time:21025ms step_avg:37.15ms
step:567/1555 train_time:21082ms step_avg:37.18ms
step:568/1555 train_time:21147ms step_avg:37.23ms
step:569/1555 train_time:21205ms step_avg:37.27ms
step:570/1555 train_time:21270ms step_avg:37.32ms
step:571/1555 train_time:21328ms step_avg:37.35ms
step:572/1555 train_time:21392ms step_avg:37.40ms
step:573/1555 train_time:21451ms step_avg:37.44ms
step:574/1555 train_time:21514ms step_avg:37.48ms
step:575/1555 train_time:21572ms step_avg:37.52ms
step:576/1555 train_time:21637ms step_avg:37.56ms
step:577/1555 train_time:21694ms step_avg:37.60ms
step:578/1555 train_time:21758ms step_avg:37.64ms
step:579/1555 train_time:21815ms step_avg:37.68ms
step:580/1555 train_time:21879ms step_avg:37.72ms
step:581/1555 train_time:21936ms step_avg:37.76ms
step:582/1555 train_time:22000ms step_avg:37.80ms
step:583/1555 train_time:22058ms step_avg:37.84ms
step:584/1555 train_time:22123ms step_avg:37.88ms
step:585/1555 train_time:22181ms step_avg:37.92ms
step:586/1555 train_time:22247ms step_avg:37.96ms
step:587/1555 train_time:22304ms step_avg:38.00ms
step:588/1555 train_time:22369ms step_avg:38.04ms
step:589/1555 train_time:22426ms step_avg:38.08ms
step:590/1555 train_time:22491ms step_avg:38.12ms
step:591/1555 train_time:22549ms step_avg:38.15ms
step:592/1555 train_time:22613ms step_avg:38.20ms
step:593/1555 train_time:22670ms step_avg:38.23ms
step:594/1555 train_time:22735ms step_avg:38.27ms
step:595/1555 train_time:22793ms step_avg:38.31ms
step:596/1555 train_time:22857ms step_avg:38.35ms
step:597/1555 train_time:22914ms step_avg:38.38ms
step:598/1555 train_time:22977ms step_avg:38.42ms
step:599/1555 train_time:23035ms step_avg:38.46ms
step:600/1555 train_time:23100ms step_avg:38.50ms
step:601/1555 train_time:23157ms step_avg:38.53ms
step:602/1555 train_time:23222ms step_avg:38.58ms
step:603/1555 train_time:23280ms step_avg:38.61ms
step:604/1555 train_time:23345ms step_avg:38.65ms
step:605/1555 train_time:23402ms step_avg:38.68ms
step:606/1555 train_time:23468ms step_avg:38.73ms
step:607/1555 train_time:23525ms step_avg:38.76ms
step:608/1555 train_time:23590ms step_avg:38.80ms
step:609/1555 train_time:23649ms step_avg:38.83ms
step:610/1555 train_time:23714ms step_avg:38.87ms
step:611/1555 train_time:23771ms step_avg:38.91ms
step:612/1555 train_time:23835ms step_avg:38.95ms
step:613/1555 train_time:23893ms step_avg:38.98ms
step:614/1555 train_time:23957ms step_avg:39.02ms
step:615/1555 train_time:24015ms step_avg:39.05ms
step:616/1555 train_time:24078ms step_avg:39.09ms
step:617/1555 train_time:24135ms step_avg:39.12ms
step:618/1555 train_time:24199ms step_avg:39.16ms
step:619/1555 train_time:24256ms step_avg:39.19ms
step:620/1555 train_time:24322ms step_avg:39.23ms
step:621/1555 train_time:24379ms step_avg:39.26ms
step:622/1555 train_time:24444ms step_avg:39.30ms
step:623/1555 train_time:24502ms step_avg:39.33ms
step:624/1555 train_time:24568ms step_avg:39.37ms
step:625/1555 train_time:24626ms step_avg:39.40ms
step:626/1555 train_time:24690ms step_avg:39.44ms
step:627/1555 train_time:24748ms step_avg:39.47ms
step:628/1555 train_time:24813ms step_avg:39.51ms
step:629/1555 train_time:24871ms step_avg:39.54ms
step:630/1555 train_time:24935ms step_avg:39.58ms
step:631/1555 train_time:24993ms step_avg:39.61ms
step:632/1555 train_time:25057ms step_avg:39.65ms
step:633/1555 train_time:25114ms step_avg:39.67ms
step:634/1555 train_time:25177ms step_avg:39.71ms
step:635/1555 train_time:25235ms step_avg:39.74ms
step:636/1555 train_time:25298ms step_avg:39.78ms
step:637/1555 train_time:25356ms step_avg:39.81ms
step:638/1555 train_time:25420ms step_avg:39.84ms
step:639/1555 train_time:25478ms step_avg:39.87ms
step:640/1555 train_time:25542ms step_avg:39.91ms
step:641/1555 train_time:25600ms step_avg:39.94ms
step:642/1555 train_time:25666ms step_avg:39.98ms
step:643/1555 train_time:25724ms step_avg:40.01ms
step:644/1555 train_time:25788ms step_avg:40.04ms
step:645/1555 train_time:25846ms step_avg:40.07ms
step:646/1555 train_time:25910ms step_avg:40.11ms
step:647/1555 train_time:25968ms step_avg:40.14ms
step:648/1555 train_time:26033ms step_avg:40.17ms
step:649/1555 train_time:26091ms step_avg:40.20ms
step:650/1555 train_time:26155ms step_avg:40.24ms
step:651/1555 train_time:26213ms step_avg:40.27ms
step:652/1555 train_time:26276ms step_avg:40.30ms
step:653/1555 train_time:26334ms step_avg:40.33ms
step:654/1555 train_time:26398ms step_avg:40.36ms
step:655/1555 train_time:26456ms step_avg:40.39ms
step:656/1555 train_time:26519ms step_avg:40.43ms
step:657/1555 train_time:26577ms step_avg:40.45ms
step:658/1555 train_time:26641ms step_avg:40.49ms
step:659/1555 train_time:26699ms step_avg:40.51ms
step:660/1555 train_time:26765ms step_avg:40.55ms
step:661/1555 train_time:26823ms step_avg:40.58ms
step:662/1555 train_time:26888ms step_avg:40.62ms
step:663/1555 train_time:26945ms step_avg:40.64ms
step:664/1555 train_time:27010ms step_avg:40.68ms
step:665/1555 train_time:27067ms step_avg:40.70ms
step:666/1555 train_time:27133ms step_avg:40.74ms
step:667/1555 train_time:27191ms step_avg:40.77ms
step:668/1555 train_time:27255ms step_avg:40.80ms
step:669/1555 train_time:27312ms step_avg:40.83ms
step:670/1555 train_time:27376ms step_avg:40.86ms
step:671/1555 train_time:27434ms step_avg:40.89ms
step:672/1555 train_time:27499ms step_avg:40.92ms
step:673/1555 train_time:27556ms step_avg:40.95ms
step:674/1555 train_time:27620ms step_avg:40.98ms
step:675/1555 train_time:27677ms step_avg:41.00ms
step:676/1555 train_time:27741ms step_avg:41.04ms
step:677/1555 train_time:27799ms step_avg:41.06ms
step:678/1555 train_time:27864ms step_avg:41.10ms
step:679/1555 train_time:27922ms step_avg:41.12ms
step:680/1555 train_time:27987ms step_avg:41.16ms
step:681/1555 train_time:28045ms step_avg:41.18ms
step:682/1555 train_time:28109ms step_avg:41.22ms
step:683/1555 train_time:28167ms step_avg:41.24ms
step:684/1555 train_time:28232ms step_avg:41.27ms
step:685/1555 train_time:28289ms step_avg:41.30ms
step:686/1555 train_time:28355ms step_avg:41.33ms
step:687/1555 train_time:28412ms step_avg:41.36ms
step:688/1555 train_time:28476ms step_avg:41.39ms
step:689/1555 train_time:28535ms step_avg:41.41ms
step:690/1555 train_time:28599ms step_avg:41.45ms
step:691/1555 train_time:28657ms step_avg:41.47ms
step:692/1555 train_time:28720ms step_avg:41.50ms
step:693/1555 train_time:28778ms step_avg:41.53ms
step:694/1555 train_time:28842ms step_avg:41.56ms
step:695/1555 train_time:28899ms step_avg:41.58ms
step:696/1555 train_time:28965ms step_avg:41.62ms
step:697/1555 train_time:29023ms step_avg:41.64ms
step:698/1555 train_time:29087ms step_avg:41.67ms
step:699/1555 train_time:29145ms step_avg:41.70ms
step:700/1555 train_time:29209ms step_avg:41.73ms
step:701/1555 train_time:29267ms step_avg:41.75ms
step:702/1555 train_time:29330ms step_avg:41.78ms
step:703/1555 train_time:29390ms step_avg:41.81ms
step:704/1555 train_time:29455ms step_avg:41.84ms
step:705/1555 train_time:29512ms step_avg:41.86ms
step:706/1555 train_time:29576ms step_avg:41.89ms
step:707/1555 train_time:29635ms step_avg:41.92ms
step:708/1555 train_time:29698ms step_avg:41.95ms
step:709/1555 train_time:29755ms step_avg:41.97ms
step:710/1555 train_time:29819ms step_avg:42.00ms
step:711/1555 train_time:29877ms step_avg:42.02ms
step:712/1555 train_time:29941ms step_avg:42.05ms
step:713/1555 train_time:29999ms step_avg:42.07ms
step:714/1555 train_time:30064ms step_avg:42.11ms
step:715/1555 train_time:30122ms step_avg:42.13ms
step:716/1555 train_time:30187ms step_avg:42.16ms
step:717/1555 train_time:30245ms step_avg:42.18ms
step:718/1555 train_time:30309ms step_avg:42.21ms
step:719/1555 train_time:30367ms step_avg:42.23ms
step:720/1555 train_time:30431ms step_avg:42.27ms
step:721/1555 train_time:30490ms step_avg:42.29ms
step:722/1555 train_time:30554ms step_avg:42.32ms
step:723/1555 train_time:30613ms step_avg:42.34ms
step:724/1555 train_time:30676ms step_avg:42.37ms
step:725/1555 train_time:30733ms step_avg:42.39ms
step:726/1555 train_time:30797ms step_avg:42.42ms
step:727/1555 train_time:30856ms step_avg:42.44ms
step:728/1555 train_time:30919ms step_avg:42.47ms
step:729/1555 train_time:30977ms step_avg:42.49ms
step:730/1555 train_time:31043ms step_avg:42.52ms
step:731/1555 train_time:31099ms step_avg:42.54ms
step:732/1555 train_time:31165ms step_avg:42.57ms
step:733/1555 train_time:31222ms step_avg:42.60ms
step:734/1555 train_time:31287ms step_avg:42.62ms
step:735/1555 train_time:31345ms step_avg:42.65ms
step:736/1555 train_time:31409ms step_avg:42.68ms
step:737/1555 train_time:31467ms step_avg:42.70ms
step:738/1555 train_time:31531ms step_avg:42.72ms
step:739/1555 train_time:31590ms step_avg:42.75ms
step:740/1555 train_time:31653ms step_avg:42.77ms
step:741/1555 train_time:31711ms step_avg:42.80ms
step:742/1555 train_time:31775ms step_avg:42.82ms
step:743/1555 train_time:31833ms step_avg:42.84ms
step:744/1555 train_time:31897ms step_avg:42.87ms
step:745/1555 train_time:31955ms step_avg:42.89ms
step:746/1555 train_time:32019ms step_avg:42.92ms
step:747/1555 train_time:32076ms step_avg:42.94ms
step:748/1555 train_time:32141ms step_avg:42.97ms
step:749/1555 train_time:32199ms step_avg:42.99ms
step:750/1555 train_time:32263ms step_avg:43.02ms
step:750/1555 val_loss:3.8721 train_time:32346ms step_avg:43.13ms
step:751/1555 train_time:32366ms step_avg:43.10ms
step:752/1555 train_time:32387ms step_avg:43.07ms
step:753/1555 train_time:32447ms step_avg:43.09ms
step:754/1555 train_time:32515ms step_avg:43.12ms
step:755/1555 train_time:32574ms step_avg:43.14ms
step:756/1555 train_time:32638ms step_avg:43.17ms
step:757/1555 train_time:32695ms step_avg:43.19ms
step:758/1555 train_time:32759ms step_avg:43.22ms
step:759/1555 train_time:32816ms step_avg:43.24ms
step:760/1555 train_time:32878ms step_avg:43.26ms
step:761/1555 train_time:32935ms step_avg:43.28ms
step:762/1555 train_time:32998ms step_avg:43.30ms
step:763/1555 train_time:33054ms step_avg:43.32ms
step:764/1555 train_time:33117ms step_avg:43.35ms
step:765/1555 train_time:33175ms step_avg:43.37ms
step:766/1555 train_time:33237ms step_avg:43.39ms
step:767/1555 train_time:33295ms step_avg:43.41ms
step:768/1555 train_time:33360ms step_avg:43.44ms
step:769/1555 train_time:33419ms step_avg:43.46ms
step:770/1555 train_time:33488ms step_avg:43.49ms
step:771/1555 train_time:33549ms step_avg:43.51ms
step:772/1555 train_time:33614ms step_avg:43.54ms
step:773/1555 train_time:33672ms step_avg:43.56ms
step:774/1555 train_time:33735ms step_avg:43.58ms
step:775/1555 train_time:33792ms step_avg:43.60ms
step:776/1555 train_time:33856ms step_avg:43.63ms
step:777/1555 train_time:33912ms step_avg:43.64ms
step:778/1555 train_time:33976ms step_avg:43.67ms
step:779/1555 train_time:34032ms step_avg:43.69ms
step:780/1555 train_time:34095ms step_avg:43.71ms
step:781/1555 train_time:34153ms step_avg:43.73ms
step:782/1555 train_time:34216ms step_avg:43.75ms
step:783/1555 train_time:34274ms step_avg:43.77ms
step:784/1555 train_time:34337ms step_avg:43.80ms
step:785/1555 train_time:34397ms step_avg:43.82ms
step:786/1555 train_time:34463ms step_avg:43.85ms
step:787/1555 train_time:34521ms step_avg:43.86ms
step:788/1555 train_time:34586ms step_avg:43.89ms
step:789/1555 train_time:34644ms step_avg:43.91ms
step:790/1555 train_time:34708ms step_avg:43.93ms
step:791/1555 train_time:34767ms step_avg:43.95ms
step:792/1555 train_time:34830ms step_avg:43.98ms
step:793/1555 train_time:34888ms step_avg:43.99ms
step:794/1555 train_time:34952ms step_avg:44.02ms
step:795/1555 train_time:35009ms step_avg:44.04ms
step:796/1555 train_time:35073ms step_avg:44.06ms
step:797/1555 train_time:35130ms step_avg:44.08ms
step:798/1555 train_time:35194ms step_avg:44.10ms
step:799/1555 train_time:35252ms step_avg:44.12ms
step:800/1555 train_time:35316ms step_avg:44.14ms
step:801/1555 train_time:35374ms step_avg:44.16ms
step:802/1555 train_time:35438ms step_avg:44.19ms
step:803/1555 train_time:35496ms step_avg:44.20ms
step:804/1555 train_time:35561ms step_avg:44.23ms
step:805/1555 train_time:35617ms step_avg:44.25ms
step:806/1555 train_time:35682ms step_avg:44.27ms
step:807/1555 train_time:35740ms step_avg:44.29ms
step:808/1555 train_time:35804ms step_avg:44.31ms
step:809/1555 train_time:35863ms step_avg:44.33ms
step:810/1555 train_time:35927ms step_avg:44.35ms
step:811/1555 train_time:35984ms step_avg:44.37ms
step:812/1555 train_time:36048ms step_avg:44.39ms
step:813/1555 train_time:36105ms step_avg:44.41ms
step:814/1555 train_time:36170ms step_avg:44.43ms
step:815/1555 train_time:36228ms step_avg:44.45ms
step:816/1555 train_time:36292ms step_avg:44.48ms
step:817/1555 train_time:36350ms step_avg:44.49ms
step:818/1555 train_time:36415ms step_avg:44.52ms
step:819/1555 train_time:36474ms step_avg:44.53ms
step:820/1555 train_time:36537ms step_avg:44.56ms
step:821/1555 train_time:36595ms step_avg:44.57ms
step:822/1555 train_time:36659ms step_avg:44.60ms
step:823/1555 train_time:36716ms step_avg:44.61ms
step:824/1555 train_time:36780ms step_avg:44.64ms
step:825/1555 train_time:36837ms step_avg:44.65ms
step:826/1555 train_time:36902ms step_avg:44.68ms
step:827/1555 train_time:36959ms step_avg:44.69ms
step:828/1555 train_time:37024ms step_avg:44.72ms
step:829/1555 train_time:37081ms step_avg:44.73ms
step:830/1555 train_time:37145ms step_avg:44.75ms
step:831/1555 train_time:37202ms step_avg:44.77ms
step:832/1555 train_time:37266ms step_avg:44.79ms
step:833/1555 train_time:37325ms step_avg:44.81ms
step:834/1555 train_time:37389ms step_avg:44.83ms
step:835/1555 train_time:37448ms step_avg:44.85ms
step:836/1555 train_time:37512ms step_avg:44.87ms
step:837/1555 train_time:37572ms step_avg:44.89ms
step:838/1555 train_time:37635ms step_avg:44.91ms
step:839/1555 train_time:37693ms step_avg:44.93ms
step:840/1555 train_time:37758ms step_avg:44.95ms
step:841/1555 train_time:37815ms step_avg:44.96ms
step:842/1555 train_time:37880ms step_avg:44.99ms
step:843/1555 train_time:37936ms step_avg:45.00ms
step:844/1555 train_time:38001ms step_avg:45.02ms
step:845/1555 train_time:38058ms step_avg:45.04ms
step:846/1555 train_time:38122ms step_avg:45.06ms
step:847/1555 train_time:38179ms step_avg:45.08ms
step:848/1555 train_time:38243ms step_avg:45.10ms
step:849/1555 train_time:38301ms step_avg:45.11ms
step:850/1555 train_time:38366ms step_avg:45.14ms
step:851/1555 train_time:38423ms step_avg:45.15ms
step:852/1555 train_time:38489ms step_avg:45.17ms
step:853/1555 train_time:38547ms step_avg:45.19ms
step:854/1555 train_time:38612ms step_avg:45.21ms
step:855/1555 train_time:38670ms step_avg:45.23ms
step:856/1555 train_time:38734ms step_avg:45.25ms
step:857/1555 train_time:38792ms step_avg:45.27ms
step:858/1555 train_time:38857ms step_avg:45.29ms
step:859/1555 train_time:38915ms step_avg:45.30ms
step:860/1555 train_time:38978ms step_avg:45.32ms
step:861/1555 train_time:39035ms step_avg:45.34ms
step:862/1555 train_time:39099ms step_avg:45.36ms
step:863/1555 train_time:39156ms step_avg:45.37ms
step:864/1555 train_time:39220ms step_avg:45.39ms
step:865/1555 train_time:39277ms step_avg:45.41ms
step:866/1555 train_time:39341ms step_avg:45.43ms
step:867/1555 train_time:39399ms step_avg:45.44ms
step:868/1555 train_time:39464ms step_avg:45.47ms
step:869/1555 train_time:39522ms step_avg:45.48ms
step:870/1555 train_time:39587ms step_avg:45.50ms
step:871/1555 train_time:39645ms step_avg:45.52ms
step:872/1555 train_time:39709ms step_avg:45.54ms
step:873/1555 train_time:39768ms step_avg:45.55ms
step:874/1555 train_time:39833ms step_avg:45.58ms
step:875/1555 train_time:39890ms step_avg:45.59ms
step:876/1555 train_time:39955ms step_avg:45.61ms
step:877/1555 train_time:40013ms step_avg:45.62ms
step:878/1555 train_time:40077ms step_avg:45.65ms
step:879/1555 train_time:40134ms step_avg:45.66ms
step:880/1555 train_time:40198ms step_avg:45.68ms
step:881/1555 train_time:40256ms step_avg:45.69ms
step:882/1555 train_time:40319ms step_avg:45.71ms
step:883/1555 train_time:40377ms step_avg:45.73ms
step:884/1555 train_time:40440ms step_avg:45.75ms
step:885/1555 train_time:40498ms step_avg:45.76ms
step:886/1555 train_time:40563ms step_avg:45.78ms
step:887/1555 train_time:40620ms step_avg:45.80ms
step:888/1555 train_time:40686ms step_avg:45.82ms
step:889/1555 train_time:40744ms step_avg:45.83ms
step:890/1555 train_time:40809ms step_avg:45.85ms
step:891/1555 train_time:40867ms step_avg:45.87ms
step:892/1555 train_time:40931ms step_avg:45.89ms
step:893/1555 train_time:40989ms step_avg:45.90ms
step:894/1555 train_time:41053ms step_avg:45.92ms
step:895/1555 train_time:41111ms step_avg:45.93ms
step:896/1555 train_time:41176ms step_avg:45.96ms
step:897/1555 train_time:41233ms step_avg:45.97ms
step:898/1555 train_time:41297ms step_avg:45.99ms
step:899/1555 train_time:41355ms step_avg:46.00ms
step:900/1555 train_time:41419ms step_avg:46.02ms
step:901/1555 train_time:41477ms step_avg:46.03ms
step:902/1555 train_time:41541ms step_avg:46.05ms
step:903/1555 train_time:41598ms step_avg:46.07ms
step:904/1555 train_time:41663ms step_avg:46.09ms
step:905/1555 train_time:41720ms step_avg:46.10ms
step:906/1555 train_time:41785ms step_avg:46.12ms
step:907/1555 train_time:41844ms step_avg:46.13ms
step:908/1555 train_time:41909ms step_avg:46.16ms
step:909/1555 train_time:41967ms step_avg:46.17ms
step:910/1555 train_time:42030ms step_avg:46.19ms
step:911/1555 train_time:42089ms step_avg:46.20ms
step:912/1555 train_time:42153ms step_avg:46.22ms
step:913/1555 train_time:42211ms step_avg:46.23ms
step:914/1555 train_time:42277ms step_avg:46.26ms
step:915/1555 train_time:42335ms step_avg:46.27ms
step:916/1555 train_time:42398ms step_avg:46.29ms
step:917/1555 train_time:42456ms step_avg:46.30ms
step:918/1555 train_time:42519ms step_avg:46.32ms
step:919/1555 train_time:42577ms step_avg:46.33ms
step:920/1555 train_time:42640ms step_avg:46.35ms
step:921/1555 train_time:42698ms step_avg:46.36ms
step:922/1555 train_time:42762ms step_avg:46.38ms
step:923/1555 train_time:42821ms step_avg:46.39ms
step:924/1555 train_time:42885ms step_avg:46.41ms
step:925/1555 train_time:42943ms step_avg:46.42ms
step:926/1555 train_time:43008ms step_avg:46.44ms
step:927/1555 train_time:43065ms step_avg:46.46ms
step:928/1555 train_time:43129ms step_avg:46.47ms
step:929/1555 train_time:43187ms step_avg:46.49ms
step:930/1555 train_time:43251ms step_avg:46.51ms
step:931/1555 train_time:43309ms step_avg:46.52ms
step:932/1555 train_time:43375ms step_avg:46.54ms
step:933/1555 train_time:43433ms step_avg:46.55ms
step:934/1555 train_time:43497ms step_avg:46.57ms
step:935/1555 train_time:43555ms step_avg:46.58ms
step:936/1555 train_time:43619ms step_avg:46.60ms
step:937/1555 train_time:43676ms step_avg:46.61ms
step:938/1555 train_time:43740ms step_avg:46.63ms
step:939/1555 train_time:43798ms step_avg:46.64ms
step:940/1555 train_time:43861ms step_avg:46.66ms
step:941/1555 train_time:43919ms step_avg:46.67ms
step:942/1555 train_time:43984ms step_avg:46.69ms
step:943/1555 train_time:44041ms step_avg:46.70ms
step:944/1555 train_time:44107ms step_avg:46.72ms
step:945/1555 train_time:44165ms step_avg:46.74ms
step:946/1555 train_time:44229ms step_avg:46.75ms
step:947/1555 train_time:44287ms step_avg:46.77ms
step:948/1555 train_time:44351ms step_avg:46.78ms
step:949/1555 train_time:44409ms step_avg:46.80ms
step:950/1555 train_time:44474ms step_avg:46.82ms
step:951/1555 train_time:44532ms step_avg:46.83ms
step:952/1555 train_time:44599ms step_avg:46.85ms
step:953/1555 train_time:44656ms step_avg:46.86ms
step:954/1555 train_time:44719ms step_avg:46.88ms
step:955/1555 train_time:44777ms step_avg:46.89ms
step:956/1555 train_time:44842ms step_avg:46.91ms
step:957/1555 train_time:44900ms step_avg:46.92ms
step:958/1555 train_time:44965ms step_avg:46.94ms
step:959/1555 train_time:45022ms step_avg:46.95ms
step:960/1555 train_time:45087ms step_avg:46.97ms
step:961/1555 train_time:45144ms step_avg:46.98ms
step:962/1555 train_time:45208ms step_avg:46.99ms
step:963/1555 train_time:45266ms step_avg:47.01ms
step:964/1555 train_time:45331ms step_avg:47.02ms
step:965/1555 train_time:45389ms step_avg:47.03ms
step:966/1555 train_time:45453ms step_avg:47.05ms
step:967/1555 train_time:45511ms step_avg:47.06ms
step:968/1555 train_time:45577ms step_avg:47.08ms
step:969/1555 train_time:45634ms step_avg:47.09ms
step:970/1555 train_time:45698ms step_avg:47.11ms
step:971/1555 train_time:45756ms step_avg:47.12ms
step:972/1555 train_time:45819ms step_avg:47.14ms
step:973/1555 train_time:45876ms step_avg:47.15ms
step:974/1555 train_time:45941ms step_avg:47.17ms
step:975/1555 train_time:45998ms step_avg:47.18ms
step:976/1555 train_time:46062ms step_avg:47.19ms
step:977/1555 train_time:46119ms step_avg:47.20ms
step:978/1555 train_time:46184ms step_avg:47.22ms
step:979/1555 train_time:46242ms step_avg:47.23ms
step:980/1555 train_time:46307ms step_avg:47.25ms
step:981/1555 train_time:46365ms step_avg:47.26ms
step:982/1555 train_time:46429ms step_avg:47.28ms
step:983/1555 train_time:46488ms step_avg:47.29ms
step:984/1555 train_time:46553ms step_avg:47.31ms
step:985/1555 train_time:46611ms step_avg:47.32ms
step:986/1555 train_time:46676ms step_avg:47.34ms
step:987/1555 train_time:46734ms step_avg:47.35ms
step:988/1555 train_time:46798ms step_avg:47.37ms
step:989/1555 train_time:46856ms step_avg:47.38ms
step:990/1555 train_time:46919ms step_avg:47.39ms
step:991/1555 train_time:46977ms step_avg:47.40ms
step:992/1555 train_time:47041ms step_avg:47.42ms
step:993/1555 train_time:47098ms step_avg:47.43ms
step:994/1555 train_time:47163ms step_avg:47.45ms
step:995/1555 train_time:47219ms step_avg:47.46ms
step:996/1555 train_time:47285ms step_avg:47.47ms
step:997/1555 train_time:47343ms step_avg:47.49ms
step:998/1555 train_time:47409ms step_avg:47.50ms
step:999/1555 train_time:47467ms step_avg:47.51ms
step:1000/1555 train_time:47530ms step_avg:47.53ms
step:1000/1555 val_loss:3.5717 train_time:47613ms step_avg:47.61ms
step:1001/1555 train_time:47631ms step_avg:47.58ms
step:1002/1555 train_time:47654ms step_avg:47.56ms
step:1003/1555 train_time:47710ms step_avg:47.57ms
step:1004/1555 train_time:47780ms step_avg:47.59ms
step:1005/1555 train_time:47838ms step_avg:47.60ms
step:1006/1555 train_time:47902ms step_avg:47.62ms
step:1007/1555 train_time:47959ms step_avg:47.63ms
step:1008/1555 train_time:48023ms step_avg:47.64ms
step:1009/1555 train_time:48080ms step_avg:47.65ms
step:1010/1555 train_time:48144ms step_avg:47.67ms
step:1011/1555 train_time:48204ms step_avg:47.68ms
step:1012/1555 train_time:48289ms step_avg:47.72ms
step:1013/1555 train_time:48373ms step_avg:47.75ms
step:1014/1555 train_time:48464ms step_avg:47.79ms
step:1015/1555 train_time:48547ms step_avg:47.83ms
step:1016/1555 train_time:48638ms step_avg:47.87ms
step:1017/1555 train_time:48724ms step_avg:47.91ms
step:1018/1555 train_time:48816ms step_avg:47.95ms
step:1019/1555 train_time:48901ms step_avg:47.99ms
step:1020/1555 train_time:48992ms step_avg:48.03ms
step:1021/1555 train_time:49076ms step_avg:48.07ms
step:1022/1555 train_time:49166ms step_avg:48.11ms
step:1023/1555 train_time:49248ms step_avg:48.14ms
step:1024/1555 train_time:49338ms step_avg:48.18ms
step:1025/1555 train_time:49421ms step_avg:48.22ms
step:1026/1555 train_time:49510ms step_avg:48.26ms
step:1027/1555 train_time:49595ms step_avg:48.29ms
step:1028/1555 train_time:49685ms step_avg:48.33ms
step:1029/1555 train_time:49771ms step_avg:48.37ms
step:1030/1555 train_time:49862ms step_avg:48.41ms
step:1031/1555 train_time:49946ms step_avg:48.44ms
step:1032/1555 train_time:50039ms step_avg:48.49ms
step:1033/1555 train_time:50122ms step_avg:48.52ms
step:1034/1555 train_time:50210ms step_avg:48.56ms
step:1035/1555 train_time:50293ms step_avg:48.59ms
step:1036/1555 train_time:50383ms step_avg:48.63ms
step:1037/1555 train_time:50466ms step_avg:48.66ms
step:1038/1555 train_time:50556ms step_avg:48.71ms
step:1039/1555 train_time:50641ms step_avg:48.74ms
step:1040/1555 train_time:50731ms step_avg:48.78ms
step:1041/1555 train_time:50817ms step_avg:48.82ms
step:1042/1555 train_time:50907ms step_avg:48.85ms
step:1043/1555 train_time:50991ms step_avg:48.89ms
step:1044/1555 train_time:51081ms step_avg:48.93ms
step:1045/1555 train_time:51165ms step_avg:48.96ms
step:1046/1555 train_time:51254ms step_avg:49.00ms
step:1047/1555 train_time:51338ms step_avg:49.03ms
step:1048/1555 train_time:51426ms step_avg:49.07ms
step:1049/1555 train_time:51510ms step_avg:49.10ms
step:1050/1555 train_time:51601ms step_avg:49.14ms
step:1051/1555 train_time:51684ms step_avg:49.18ms
step:1052/1555 train_time:51776ms step_avg:49.22ms
step:1053/1555 train_time:51861ms step_avg:49.25ms
step:1054/1555 train_time:51951ms step_avg:49.29ms
step:1055/1555 train_time:52036ms step_avg:49.32ms
step:1056/1555 train_time:52127ms step_avg:49.36ms
step:1057/1555 train_time:52211ms step_avg:49.40ms
step:1058/1555 train_time:52301ms step_avg:49.43ms
step:1059/1555 train_time:52384ms step_avg:49.47ms
step:1060/1555 train_time:52473ms step_avg:49.50ms
step:1061/1555 train_time:52557ms step_avg:49.54ms
step:1062/1555 train_time:52648ms step_avg:49.57ms
step:1063/1555 train_time:52733ms step_avg:49.61ms
step:1064/1555 train_time:52823ms step_avg:49.65ms
step:1065/1555 train_time:52908ms step_avg:49.68ms
step:1066/1555 train_time:52999ms step_avg:49.72ms
step:1067/1555 train_time:53083ms step_avg:49.75ms
step:1068/1555 train_time:53172ms step_avg:49.79ms
step:1069/1555 train_time:53257ms step_avg:49.82ms
step:1070/1555 train_time:53346ms step_avg:49.86ms
step:1071/1555 train_time:53430ms step_avg:49.89ms
step:1072/1555 train_time:53521ms step_avg:49.93ms
step:1073/1555 train_time:53604ms step_avg:49.96ms
step:1074/1555 train_time:53695ms step_avg:50.00ms
step:1075/1555 train_time:53779ms step_avg:50.03ms
step:1076/1555 train_time:53868ms step_avg:50.06ms
step:1077/1555 train_time:53952ms step_avg:50.09ms
step:1078/1555 train_time:54043ms step_avg:50.13ms
step:1079/1555 train_time:54127ms step_avg:50.16ms
step:1080/1555 train_time:54219ms step_avg:50.20ms
step:1081/1555 train_time:54301ms step_avg:50.23ms
step:1082/1555 train_time:54391ms step_avg:50.27ms
step:1083/1555 train_time:54475ms step_avg:50.30ms
step:1084/1555 train_time:54565ms step_avg:50.34ms
step:1085/1555 train_time:54649ms step_avg:50.37ms
step:1086/1555 train_time:54739ms step_avg:50.40ms
step:1087/1555 train_time:54823ms step_avg:50.44ms
step:1088/1555 train_time:54914ms step_avg:50.47ms
step:1089/1555 train_time:54998ms step_avg:50.50ms
step:1090/1555 train_time:55089ms step_avg:50.54ms
step:1091/1555 train_time:55173ms step_avg:50.57ms
step:1092/1555 train_time:55263ms step_avg:50.61ms
step:1093/1555 train_time:55346ms step_avg:50.64ms
step:1094/1555 train_time:55436ms step_avg:50.67ms
step:1095/1555 train_time:55520ms step_avg:50.70ms
step:1096/1555 train_time:55610ms step_avg:50.74ms
step:1097/1555 train_time:55694ms step_avg:50.77ms
step:1098/1555 train_time:55784ms step_avg:50.81ms
step:1099/1555 train_time:55868ms step_avg:50.84ms
step:1100/1555 train_time:55959ms step_avg:50.87ms
step:1101/1555 train_time:56043ms step_avg:50.90ms
step:1102/1555 train_time:56134ms step_avg:50.94ms
step:1103/1555 train_time:56219ms step_avg:50.97ms
step:1104/1555 train_time:56308ms step_avg:51.00ms
step:1105/1555 train_time:56391ms step_avg:51.03ms
step:1106/1555 train_time:56482ms step_avg:51.07ms
step:1107/1555 train_time:56565ms step_avg:51.10ms
step:1108/1555 train_time:56655ms step_avg:51.13ms
step:1109/1555 train_time:56739ms step_avg:51.16ms
step:1110/1555 train_time:56829ms step_avg:51.20ms
step:1111/1555 train_time:56914ms step_avg:51.23ms
step:1112/1555 train_time:57003ms step_avg:51.26ms
step:1113/1555 train_time:57087ms step_avg:51.29ms
step:1114/1555 train_time:57177ms step_avg:51.33ms
step:1115/1555 train_time:57261ms step_avg:51.35ms
step:1116/1555 train_time:57350ms step_avg:51.39ms
step:1117/1555 train_time:57435ms step_avg:51.42ms
step:1118/1555 train_time:57525ms step_avg:51.45ms
step:1119/1555 train_time:57608ms step_avg:51.48ms
step:1120/1555 train_time:57699ms step_avg:51.52ms
step:1121/1555 train_time:57783ms step_avg:51.55ms
step:1122/1555 train_time:57873ms step_avg:51.58ms
step:1123/1555 train_time:57958ms step_avg:51.61ms
step:1124/1555 train_time:58048ms step_avg:51.64ms
step:1125/1555 train_time:58133ms step_avg:51.67ms
step:1126/1555 train_time:58223ms step_avg:51.71ms
step:1127/1555 train_time:58306ms step_avg:51.74ms
step:1128/1555 train_time:58396ms step_avg:51.77ms
step:1129/1555 train_time:58480ms step_avg:51.80ms
step:1130/1555 train_time:58569ms step_avg:51.83ms
step:1131/1555 train_time:58654ms step_avg:51.86ms
step:1132/1555 train_time:58744ms step_avg:51.89ms
step:1133/1555 train_time:58828ms step_avg:51.92ms
step:1134/1555 train_time:58920ms step_avg:51.96ms
step:1135/1555 train_time:59003ms step_avg:51.98ms
step:1136/1555 train_time:59093ms step_avg:52.02ms
step:1137/1555 train_time:59177ms step_avg:52.05ms
step:1138/1555 train_time:59266ms step_avg:52.08ms
step:1139/1555 train_time:59352ms step_avg:52.11ms
step:1140/1555 train_time:59441ms step_avg:52.14ms
step:1141/1555 train_time:59525ms step_avg:52.17ms
step:1142/1555 train_time:59616ms step_avg:52.20ms
step:1143/1555 train_time:59700ms step_avg:52.23ms
step:1144/1555 train_time:59790ms step_avg:52.26ms
step:1145/1555 train_time:59874ms step_avg:52.29ms
step:1146/1555 train_time:59964ms step_avg:52.32ms
step:1147/1555 train_time:60048ms step_avg:52.35ms
step:1148/1555 train_time:60139ms step_avg:52.39ms
step:1149/1555 train_time:60222ms step_avg:52.41ms
step:1150/1555 train_time:60313ms step_avg:52.45ms
step:1151/1555 train_time:60397ms step_avg:52.47ms
step:1152/1555 train_time:60486ms step_avg:52.51ms
step:1153/1555 train_time:60571ms step_avg:52.53ms
step:1154/1555 train_time:60661ms step_avg:52.57ms
step:1155/1555 train_time:60745ms step_avg:52.59ms
step:1156/1555 train_time:60835ms step_avg:52.63ms
step:1157/1555 train_time:60919ms step_avg:52.65ms
step:1158/1555 train_time:61008ms step_avg:52.68ms
step:1159/1555 train_time:61093ms step_avg:52.71ms
step:1160/1555 train_time:61182ms step_avg:52.74ms
step:1161/1555 train_time:61266ms step_avg:52.77ms
step:1162/1555 train_time:61357ms step_avg:52.80ms
step:1163/1555 train_time:61440ms step_avg:52.83ms
step:1164/1555 train_time:61530ms step_avg:52.86ms
step:1165/1555 train_time:61615ms step_avg:52.89ms
step:1166/1555 train_time:61705ms step_avg:52.92ms
step:1167/1555 train_time:61789ms step_avg:52.95ms
step:1168/1555 train_time:61879ms step_avg:52.98ms
step:1169/1555 train_time:61962ms step_avg:53.00ms
step:1170/1555 train_time:62053ms step_avg:53.04ms
step:1171/1555 train_time:62137ms step_avg:53.06ms
step:1172/1555 train_time:62226ms step_avg:53.09ms
step:1173/1555 train_time:62312ms step_avg:53.12ms
step:1174/1555 train_time:62402ms step_avg:53.15ms
step:1175/1555 train_time:62485ms step_avg:53.18ms
step:1176/1555 train_time:62576ms step_avg:53.21ms
step:1177/1555 train_time:62659ms step_avg:53.24ms
step:1178/1555 train_time:62749ms step_avg:53.27ms
step:1179/1555 train_time:62834ms step_avg:53.29ms
step:1180/1555 train_time:62924ms step_avg:53.33ms
step:1181/1555 train_time:63008ms step_avg:53.35ms
step:1182/1555 train_time:63099ms step_avg:53.38ms
step:1183/1555 train_time:63182ms step_avg:53.41ms
step:1184/1555 train_time:63271ms step_avg:53.44ms
step:1185/1555 train_time:63357ms step_avg:53.47ms
step:1186/1555 train_time:63446ms step_avg:53.50ms
step:1187/1555 train_time:63530ms step_avg:53.52ms
step:1188/1555 train_time:63621ms step_avg:53.55ms
step:1189/1555 train_time:63703ms step_avg:53.58ms
step:1190/1555 train_time:63793ms step_avg:53.61ms
step:1191/1555 train_time:63877ms step_avg:53.63ms
step:1192/1555 train_time:63966ms step_avg:53.66ms
step:1193/1555 train_time:64051ms step_avg:53.69ms
step:1194/1555 train_time:64142ms step_avg:53.72ms
step:1195/1555 train_time:64225ms step_avg:53.75ms
step:1196/1555 train_time:64317ms step_avg:53.78ms
step:1197/1555 train_time:64400ms step_avg:53.80ms
step:1198/1555 train_time:64490ms step_avg:53.83ms
step:1199/1555 train_time:64574ms step_avg:53.86ms
step:1200/1555 train_time:64665ms step_avg:53.89ms
step:1201/1555 train_time:64749ms step_avg:53.91ms
step:1202/1555 train_time:64840ms step_avg:53.94ms
step:1203/1555 train_time:64924ms step_avg:53.97ms
step:1204/1555 train_time:65014ms step_avg:54.00ms
step:1205/1555 train_time:65098ms step_avg:54.02ms
step:1206/1555 train_time:65186ms step_avg:54.05ms
step:1207/1555 train_time:65271ms step_avg:54.08ms
step:1208/1555 train_time:65361ms step_avg:54.11ms
step:1209/1555 train_time:65445ms step_avg:54.13ms
step:1210/1555 train_time:65536ms step_avg:54.16ms
step:1211/1555 train_time:65620ms step_avg:54.19ms
step:1212/1555 train_time:65710ms step_avg:54.22ms
step:1213/1555 train_time:65795ms step_avg:54.24ms
step:1214/1555 train_time:65884ms step_avg:54.27ms
step:1215/1555 train_time:65968ms step_avg:54.29ms
step:1216/1555 train_time:66060ms step_avg:54.33ms
step:1217/1555 train_time:66143ms step_avg:54.35ms
step:1218/1555 train_time:66233ms step_avg:54.38ms
step:1219/1555 train_time:66317ms step_avg:54.40ms
step:1220/1555 train_time:66405ms step_avg:54.43ms
step:1221/1555 train_time:66489ms step_avg:54.45ms
step:1222/1555 train_time:66581ms step_avg:54.49ms
step:1223/1555 train_time:66664ms step_avg:54.51ms
step:1224/1555 train_time:66754ms step_avg:54.54ms
step:1225/1555 train_time:66839ms step_avg:54.56ms
step:1226/1555 train_time:66928ms step_avg:54.59ms
step:1227/1555 train_time:67015ms step_avg:54.62ms
step:1228/1555 train_time:67103ms step_avg:54.64ms
step:1229/1555 train_time:67188ms step_avg:54.67ms
step:1230/1555 train_time:67278ms step_avg:54.70ms
step:1231/1555 train_time:67362ms step_avg:54.72ms
step:1232/1555 train_time:67451ms step_avg:54.75ms
step:1233/1555 train_time:67535ms step_avg:54.77ms
step:1234/1555 train_time:67626ms step_avg:54.80ms
step:1235/1555 train_time:67710ms step_avg:54.83ms
step:1236/1555 train_time:67800ms step_avg:54.85ms
step:1237/1555 train_time:67884ms step_avg:54.88ms
step:1238/1555 train_time:67974ms step_avg:54.91ms
step:1239/1555 train_time:68060ms step_avg:54.93ms
step:1240/1555 train_time:68148ms step_avg:54.96ms
step:1241/1555 train_time:68233ms step_avg:54.98ms
step:1242/1555 train_time:68323ms step_avg:55.01ms
step:1243/1555 train_time:68406ms step_avg:55.03ms
step:1244/1555 train_time:68496ms step_avg:55.06ms
step:1245/1555 train_time:68580ms step_avg:55.08ms
step:1246/1555 train_time:68670ms step_avg:55.11ms
step:1247/1555 train_time:68754ms step_avg:55.14ms
step:1248/1555 train_time:68844ms step_avg:55.16ms
step:1249/1555 train_time:68928ms step_avg:55.19ms
step:1250/1555 train_time:69020ms step_avg:55.22ms
step:1250/1555 val_loss:3.3999 train_time:69133ms step_avg:55.31ms
step:1251/1555 train_time:69151ms step_avg:55.28ms
step:1252/1555 train_time:69195ms step_avg:55.27ms
step:1253/1555 train_time:69283ms step_avg:55.29ms
step:1254/1555 train_time:69373ms step_avg:55.32ms
step:1255/1555 train_time:69456ms step_avg:55.34ms
step:1256/1555 train_time:69547ms step_avg:55.37ms
step:1257/1555 train_time:69629ms step_avg:55.39ms
step:1258/1555 train_time:69718ms step_avg:55.42ms
step:1259/1555 train_time:69801ms step_avg:55.44ms
step:1260/1555 train_time:69891ms step_avg:55.47ms
step:1261/1555 train_time:69973ms step_avg:55.49ms
step:1262/1555 train_time:70063ms step_avg:55.52ms
step:1263/1555 train_time:70150ms step_avg:55.54ms
step:1264/1555 train_time:70242ms step_avg:55.57ms
step:1265/1555 train_time:70330ms step_avg:55.60ms
step:1266/1555 train_time:70421ms step_avg:55.63ms
step:1267/1555 train_time:70506ms step_avg:55.65ms
step:1268/1555 train_time:70596ms step_avg:55.67ms
step:1269/1555 train_time:70678ms step_avg:55.70ms
step:1270/1555 train_time:70768ms step_avg:55.72ms
step:1271/1555 train_time:70851ms step_avg:55.74ms
step:1272/1555 train_time:70940ms step_avg:55.77ms
step:1273/1555 train_time:71025ms step_avg:55.79ms
step:1274/1555 train_time:71115ms step_avg:55.82ms
step:1275/1555 train_time:71200ms step_avg:55.84ms
step:1276/1555 train_time:71292ms step_avg:55.87ms
step:1277/1555 train_time:71376ms step_avg:55.89ms
step:1278/1555 train_time:71468ms step_avg:55.92ms
step:1279/1555 train_time:71552ms step_avg:55.94ms
step:1280/1555 train_time:71641ms step_avg:55.97ms
step:1281/1555 train_time:71725ms step_avg:55.99ms
step:1282/1555 train_time:71814ms step_avg:56.02ms
step:1283/1555 train_time:71897ms step_avg:56.04ms
step:1284/1555 train_time:71986ms step_avg:56.06ms
step:1285/1555 train_time:72071ms step_avg:56.09ms
step:1286/1555 train_time:72161ms step_avg:56.11ms
step:1287/1555 train_time:72247ms step_avg:56.14ms
step:1288/1555 train_time:72337ms step_avg:56.16ms
step:1289/1555 train_time:72423ms step_avg:56.19ms
step:1290/1555 train_time:72514ms step_avg:56.21ms
step:1291/1555 train_time:72598ms step_avg:56.23ms
step:1292/1555 train_time:72688ms step_avg:56.26ms
step:1293/1555 train_time:72771ms step_avg:56.28ms
step:1294/1555 train_time:72860ms step_avg:56.31ms
step:1295/1555 train_time:72944ms step_avg:56.33ms
step:1296/1555 train_time:73034ms step_avg:56.35ms
step:1297/1555 train_time:73118ms step_avg:56.37ms
step:1298/1555 train_time:73210ms step_avg:56.40ms
step:1299/1555 train_time:73294ms step_avg:56.42ms
step:1300/1555 train_time:73384ms step_avg:56.45ms
step:1301/1555 train_time:73468ms step_avg:56.47ms
step:1302/1555 train_time:73559ms step_avg:56.50ms
step:1303/1555 train_time:73644ms step_avg:56.52ms
step:1304/1555 train_time:73735ms step_avg:56.54ms
step:1305/1555 train_time:73817ms step_avg:56.56ms
step:1306/1555 train_time:73907ms step_avg:56.59ms
step:1307/1555 train_time:73991ms step_avg:56.61ms
step:1308/1555 train_time:74080ms step_avg:56.64ms
step:1309/1555 train_time:74165ms step_avg:56.66ms
step:1310/1555 train_time:74255ms step_avg:56.68ms
step:1311/1555 train_time:74339ms step_avg:56.70ms
step:1312/1555 train_time:74432ms step_avg:56.73ms
step:1313/1555 train_time:74515ms step_avg:56.75ms
step:1314/1555 train_time:74605ms step_avg:56.78ms
step:1315/1555 train_time:74688ms step_avg:56.80ms
step:1316/1555 train_time:74778ms step_avg:56.82ms
step:1317/1555 train_time:74861ms step_avg:56.84ms
step:1318/1555 train_time:74951ms step_avg:56.87ms
step:1319/1555 train_time:75035ms step_avg:56.89ms
step:1320/1555 train_time:75125ms step_avg:56.91ms
step:1321/1555 train_time:75208ms step_avg:56.93ms
step:1322/1555 train_time:75299ms step_avg:56.96ms
step:1323/1555 train_time:75383ms step_avg:56.98ms
step:1324/1555 train_time:75474ms step_avg:57.00ms
step:1325/1555 train_time:75558ms step_avg:57.02ms
step:1326/1555 train_time:75649ms step_avg:57.05ms
step:1327/1555 train_time:75734ms step_avg:57.07ms
step:1328/1555 train_time:75823ms step_avg:57.10ms
step:1329/1555 train_time:75906ms step_avg:57.11ms
step:1330/1555 train_time:75995ms step_avg:57.14ms
step:1331/1555 train_time:76078ms step_avg:57.16ms
step:1332/1555 train_time:76171ms step_avg:57.19ms
step:1333/1555 train_time:76254ms step_avg:57.21ms
step:1334/1555 train_time:76344ms step_avg:57.23ms
step:1335/1555 train_time:76429ms step_avg:57.25ms
step:1336/1555 train_time:76518ms step_avg:57.27ms
step:1337/1555 train_time:76604ms step_avg:57.30ms
step:1338/1555 train_time:76694ms step_avg:57.32ms
step:1339/1555 train_time:76777ms step_avg:57.34ms
step:1340/1555 train_time:76868ms step_avg:57.36ms
step:1341/1555 train_time:76953ms step_avg:57.38ms
step:1342/1555 train_time:77041ms step_avg:57.41ms
step:1343/1555 train_time:77127ms step_avg:57.43ms
step:1344/1555 train_time:77217ms step_avg:57.45ms
step:1345/1555 train_time:77302ms step_avg:57.47ms
step:1346/1555 train_time:77393ms step_avg:57.50ms
step:1347/1555 train_time:77477ms step_avg:57.52ms
step:1348/1555 train_time:77567ms step_avg:57.54ms
step:1349/1555 train_time:77651ms step_avg:57.56ms
step:1350/1555 train_time:77740ms step_avg:57.59ms
step:1351/1555 train_time:77825ms step_avg:57.61ms
step:1352/1555 train_time:77915ms step_avg:57.63ms
step:1353/1555 train_time:77999ms step_avg:57.65ms
step:1354/1555 train_time:78089ms step_avg:57.67ms
step:1355/1555 train_time:78173ms step_avg:57.69ms
step:1356/1555 train_time:78263ms step_avg:57.72ms
step:1357/1555 train_time:78348ms step_avg:57.74ms
step:1358/1555 train_time:78438ms step_avg:57.76ms
step:1359/1555 train_time:78523ms step_avg:57.78ms
step:1360/1555 train_time:78613ms step_avg:57.80ms
step:1361/1555 train_time:78696ms step_avg:57.82ms
step:1362/1555 train_time:78787ms step_avg:57.85ms
step:1363/1555 train_time:78871ms step_avg:57.87ms
step:1364/1555 train_time:78961ms step_avg:57.89ms
step:1365/1555 train_time:79045ms step_avg:57.91ms
step:1366/1555 train_time:79135ms step_avg:57.93ms
step:1367/1555 train_time:79219ms step_avg:57.95ms
step:1368/1555 train_time:79308ms step_avg:57.97ms
step:1369/1555 train_time:79393ms step_avg:57.99ms
step:1370/1555 train_time:79483ms step_avg:58.02ms
step:1371/1555 train_time:79568ms step_avg:58.04ms
step:1372/1555 train_time:79657ms step_avg:58.06ms
step:1373/1555 train_time:79742ms step_avg:58.08ms
step:1374/1555 train_time:79833ms step_avg:58.10ms
step:1375/1555 train_time:79916ms step_avg:58.12ms
step:1376/1555 train_time:80007ms step_avg:58.14ms
step:1377/1555 train_time:80091ms step_avg:58.16ms
step:1378/1555 train_time:80180ms step_avg:58.19ms
step:1379/1555 train_time:80264ms step_avg:58.20ms
step:1380/1555 train_time:80355ms step_avg:58.23ms
step:1381/1555 train_time:80438ms step_avg:58.25ms
step:1382/1555 train_time:80529ms step_avg:58.27ms
step:1383/1555 train_time:80612ms step_avg:58.29ms
step:1384/1555 train_time:80702ms step_avg:58.31ms
step:1385/1555 train_time:80786ms step_avg:58.33ms
step:1386/1555 train_time:80875ms step_avg:58.35ms
step:1387/1555 train_time:80960ms step_avg:58.37ms
step:1388/1555 train_time:81050ms step_avg:58.39ms
step:1389/1555 train_time:81134ms step_avg:58.41ms
step:1390/1555 train_time:81224ms step_avg:58.43ms
step:1391/1555 train_time:81308ms step_avg:58.45ms
step:1392/1555 train_time:81398ms step_avg:58.48ms
step:1393/1555 train_time:81482ms step_avg:58.49ms
step:1394/1555 train_time:81572ms step_avg:58.52ms
step:1395/1555 train_time:81655ms step_avg:58.53ms
step:1396/1555 train_time:81746ms step_avg:58.56ms
step:1397/1555 train_time:81831ms step_avg:58.58ms
step:1398/1555 train_time:81919ms step_avg:58.60ms
step:1399/1555 train_time:82003ms step_avg:58.62ms
step:1400/1555 train_time:82093ms step_avg:58.64ms
step:1401/1555 train_time:82177ms step_avg:58.66ms
step:1402/1555 train_time:82268ms step_avg:58.68ms
step:1403/1555 train_time:82353ms step_avg:58.70ms
step:1404/1555 train_time:82443ms step_avg:58.72ms
step:1405/1555 train_time:82527ms step_avg:58.74ms
step:1406/1555 train_time:82616ms step_avg:58.76ms
step:1407/1555 train_time:82700ms step_avg:58.78ms
step:1408/1555 train_time:82792ms step_avg:58.80ms
step:1409/1555 train_time:82875ms step_avg:58.82ms
step:1410/1555 train_time:82966ms step_avg:58.84ms
step:1411/1555 train_time:83050ms step_avg:58.86ms
step:1412/1555 train_time:83139ms step_avg:58.88ms
step:1413/1555 train_time:83224ms step_avg:58.90ms
step:1414/1555 train_time:83314ms step_avg:58.92ms
step:1415/1555 train_time:83398ms step_avg:58.94ms
step:1416/1555 train_time:83487ms step_avg:58.96ms
step:1417/1555 train_time:83572ms step_avg:58.98ms
step:1418/1555 train_time:83661ms step_avg:59.00ms
step:1419/1555 train_time:83746ms step_avg:59.02ms
step:1420/1555 train_time:83836ms step_avg:59.04ms
step:1421/1555 train_time:83919ms step_avg:59.06ms
step:1422/1555 train_time:84010ms step_avg:59.08ms
step:1423/1555 train_time:84094ms step_avg:59.10ms
step:1424/1555 train_time:84184ms step_avg:59.12ms
step:1425/1555 train_time:84268ms step_avg:59.14ms
step:1426/1555 train_time:84358ms step_avg:59.16ms
step:1427/1555 train_time:84442ms step_avg:59.17ms
step:1428/1555 train_time:84532ms step_avg:59.20ms
step:1429/1555 train_time:84616ms step_avg:59.21ms
step:1430/1555 train_time:84706ms step_avg:59.24ms
step:1431/1555 train_time:84792ms step_avg:59.25ms
step:1432/1555 train_time:84881ms step_avg:59.27ms
step:1433/1555 train_time:84966ms step_avg:59.29ms
step:1434/1555 train_time:85055ms step_avg:59.31ms
step:1435/1555 train_time:85139ms step_avg:59.33ms
step:1436/1555 train_time:85230ms step_avg:59.35ms
step:1437/1555 train_time:85314ms step_avg:59.37ms
step:1438/1555 train_time:85403ms step_avg:59.39ms
step:1439/1555 train_time:85488ms step_avg:59.41ms
step:1440/1555 train_time:85578ms step_avg:59.43ms
step:1441/1555 train_time:85662ms step_avg:59.45ms
step:1442/1555 train_time:85754ms step_avg:59.47ms
step:1443/1555 train_time:85837ms step_avg:59.49ms
step:1444/1555 train_time:85928ms step_avg:59.51ms
step:1445/1555 train_time:86011ms step_avg:59.52ms
step:1446/1555 train_time:86101ms step_avg:59.54ms
step:1447/1555 train_time:86186ms step_avg:59.56ms
step:1448/1555 train_time:86276ms step_avg:59.58ms
step:1449/1555 train_time:86360ms step_avg:59.60ms
step:1450/1555 train_time:86450ms step_avg:59.62ms
step:1451/1555 train_time:86535ms step_avg:59.64ms
step:1452/1555 train_time:86625ms step_avg:59.66ms
step:1453/1555 train_time:86708ms step_avg:59.68ms
step:1454/1555 train_time:86799ms step_avg:59.70ms
step:1455/1555 train_time:86883ms step_avg:59.71ms
step:1456/1555 train_time:86974ms step_avg:59.74ms
step:1457/1555 train_time:87057ms step_avg:59.75ms
step:1458/1555 train_time:87148ms step_avg:59.77ms
step:1459/1555 train_time:87232ms step_avg:59.79ms
step:1460/1555 train_time:87321ms step_avg:59.81ms
step:1461/1555 train_time:87405ms step_avg:59.83ms
step:1462/1555 train_time:87495ms step_avg:59.85ms
step:1463/1555 train_time:87580ms step_avg:59.86ms
step:1464/1555 train_time:87670ms step_avg:59.88ms
step:1465/1555 train_time:87754ms step_avg:59.90ms
step:1466/1555 train_time:87844ms step_avg:59.92ms
step:1467/1555 train_time:87930ms step_avg:59.94ms
step:1468/1555 train_time:88018ms step_avg:59.96ms
step:1469/1555 train_time:88103ms step_avg:59.97ms
step:1470/1555 train_time:88194ms step_avg:60.00ms
step:1471/1555 train_time:88276ms step_avg:60.01ms
step:1472/1555 train_time:88366ms step_avg:60.03ms
step:1473/1555 train_time:88450ms step_avg:60.05ms
step:1474/1555 train_time:88540ms step_avg:60.07ms
step:1475/1555 train_time:88625ms step_avg:60.08ms
step:1476/1555 train_time:88715ms step_avg:60.10ms
step:1477/1555 train_time:88799ms step_avg:60.12ms
step:1478/1555 train_time:88889ms step_avg:60.14ms
step:1479/1555 train_time:88973ms step_avg:60.16ms
step:1480/1555 train_time:89063ms step_avg:60.18ms
step:1481/1555 train_time:89148ms step_avg:60.19ms
step:1482/1555 train_time:89238ms step_avg:60.21ms
step:1483/1555 train_time:89323ms step_avg:60.23ms
step:1484/1555 train_time:89414ms step_avg:60.25ms
step:1485/1555 train_time:89498ms step_avg:60.27ms
step:1486/1555 train_time:89588ms step_avg:60.29ms
step:1487/1555 train_time:89671ms step_avg:60.30ms
step:1488/1555 train_time:89761ms step_avg:60.32ms
step:1489/1555 train_time:89847ms step_avg:60.34ms
step:1490/1555 train_time:89937ms step_avg:60.36ms
step:1491/1555 train_time:90021ms step_avg:60.38ms
step:1492/1555 train_time:90111ms step_avg:60.40ms
step:1493/1555 train_time:90195ms step_avg:60.41ms
step:1494/1555 train_time:90284ms step_avg:60.43ms
step:1495/1555 train_time:90370ms step_avg:60.45ms
step:1496/1555 train_time:90459ms step_avg:60.47ms
step:1497/1555 train_time:90544ms step_avg:60.48ms
step:1498/1555 train_time:90635ms step_avg:60.50ms
step:1499/1555 train_time:90718ms step_avg:60.52ms
step:1500/1555 train_time:90808ms step_avg:60.54ms
step:1500/1555 val_loss:3.2959 train_time:90924ms step_avg:60.62ms
step:1501/1555 train_time:90942ms step_avg:60.59ms
step:1502/1555 train_time:90986ms step_avg:60.58ms
step:1503/1555 train_time:91071ms step_avg:60.59ms
step:1504/1555 train_time:91164ms step_avg:60.61ms
step:1505/1555 train_time:91250ms step_avg:60.63ms
step:1506/1555 train_time:91341ms step_avg:60.65ms
step:1507/1555 train_time:91423ms step_avg:60.67ms
step:1508/1555 train_time:91512ms step_avg:60.68ms
step:1509/1555 train_time:91595ms step_avg:60.70ms
step:1510/1555 train_time:91684ms step_avg:60.72ms
step:1511/1555 train_time:91767ms step_avg:60.73ms
step:1512/1555 train_time:91856ms step_avg:60.75ms
step:1513/1555 train_time:91943ms step_avg:60.77ms
step:1514/1555 train_time:92033ms step_avg:60.79ms
step:1515/1555 train_time:92120ms step_avg:60.81ms
step:1516/1555 train_time:92216ms step_avg:60.83ms
step:1517/1555 train_time:92302ms step_avg:60.85ms
step:1518/1555 train_time:92391ms step_avg:60.86ms
step:1519/1555 train_time:92476ms step_avg:60.88ms
step:1520/1555 train_time:92565ms step_avg:60.90ms
step:1521/1555 train_time:92648ms step_avg:60.91ms
step:1522/1555 train_time:92737ms step_avg:60.93ms
step:1523/1555 train_time:92821ms step_avg:60.95ms
step:1524/1555 train_time:92911ms step_avg:60.97ms
step:1525/1555 train_time:92996ms step_avg:60.98ms
step:1526/1555 train_time:93089ms step_avg:61.00ms
step:1527/1555 train_time:93175ms step_avg:61.02ms
step:1528/1555 train_time:93266ms step_avg:61.04ms
step:1529/1555 train_time:93350ms step_avg:61.05ms
step:1530/1555 train_time:93441ms step_avg:61.07ms
step:1531/1555 train_time:93524ms step_avg:61.09ms
step:1532/1555 train_time:93613ms step_avg:61.10ms
step:1533/1555 train_time:93696ms step_avg:61.12ms
step:1534/1555 train_time:93787ms step_avg:61.14ms
step:1535/1555 train_time:93870ms step_avg:61.15ms
step:1536/1555 train_time:93961ms step_avg:61.17ms
step:1537/1555 train_time:94047ms step_avg:61.19ms
step:1538/1555 train_time:94139ms step_avg:61.21ms
step:1539/1555 train_time:94224ms step_avg:61.22ms
step:1540/1555 train_time:94314ms step_avg:61.24ms
step:1541/1555 train_time:94399ms step_avg:61.26ms
step:1542/1555 train_time:94490ms step_avg:61.28ms
step:1543/1555 train_time:94574ms step_avg:61.29ms
step:1544/1555 train_time:94663ms step_avg:61.31ms
step:1545/1555 train_time:94748ms step_avg:61.33ms
step:1546/1555 train_time:94837ms step_avg:61.34ms
step:1547/1555 train_time:94921ms step_avg:61.36ms
step:1548/1555 train_time:95011ms step_avg:61.38ms
step:1549/1555 train_time:95096ms step_avg:61.39ms
step:1550/1555 train_time:95189ms step_avg:61.41ms
step:1551/1555 train_time:95275ms step_avg:61.43ms
step:1552/1555 train_time:95366ms step_avg:61.45ms
step:1553/1555 train_time:95450ms step_avg:61.46ms
step:1554/1555 train_time:95540ms step_avg:61.48ms
step:1555/1555 train_time:95624ms step_avg:61.49ms
step:1555/1555 val_loss:3.2797 train_time:95739ms step_avg:61.57ms
peak memory allocated: 31630 MiB reserved: 46578 MiB