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:13:14 2026       
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 570.148.08             Driver Version: 570.148.08     CUDA Version: 12.8     |
|-----------------------------------------+------------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  NVIDIA H100 80GB HBM3          On  |   00000000: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            123W /  700W |    1519MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   2  NVIDIA H100 80GB HBM3          On  |   00000000:71:00.0 Off |                    0 |
| N/A   39C    P0            125W /  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   32C    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            122W /  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           20699      C   /usr/local/bin/python                  1510MiB |
|    1   N/A  N/A           20700      C   /usr/local/bin/python                  1510MiB |
|    2   N/A  N/A           20701      C   /usr/local/bin/python                  1510MiB |
|    3   N/A  N/A           20702      C   /usr/local/bin/python                  1510MiB |
|    4   N/A  N/A           20703      C   /usr/local/bin/python                  1510MiB |
|    5   N/A  N/A           20704      C   /usr/local/bin/python                  1510MiB |
|    6   N/A  N/A           20705      C   /usr/local/bin/python                  1510MiB |
|    7   N/A  N/A           20706      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.8327 train_time:0ms step_avg:0.03ms
step:1/1555 train_time:84ms step_avg:83.57ms
step:2/1555 train_time:108ms step_avg:54.04ms
step:3/1555 train_time:129ms step_avg:42.97ms
step:4/1555 train_time:150ms step_avg:37.48ms
step:5/1555 train_time:180ms step_avg:35.94ms
step:6/1555 train_time:217ms step_avg:36.15ms
step:7/1555 train_time:248ms step_avg:35.42ms
step:8/1555 train_time:286ms step_avg:35.76ms
step:9/1555 train_time:317ms step_avg:35.21ms
step:10/1555 train_time:355ms step_avg:35.45ms
step:11/1555 train_time:385ms step_avg:35.01ms
step:12/1555 train_time:423ms step_avg:35.23ms
step:13/1555 train_time:454ms step_avg:34.89ms
step:14/1555 train_time:491ms step_avg:35.10ms
step:15/1555 train_time:522ms step_avg:34.82ms
step:16/1555 train_time:560ms step_avg:35.02ms
step:17/1555 train_time:591ms step_avg:34.78ms
step:18/1555 train_time:629ms step_avg:34.92ms
step:19/1555 train_time:660ms step_avg:34.72ms
step:20/1555 train_time:697ms step_avg:34.86ms
step:21/1555 train_time:728ms step_avg:34.69ms
step:22/1555 train_time:766ms step_avg:34.81ms
step:23/1555 train_time:797ms step_avg:34.63ms
step:24/1555 train_time:834ms step_avg:34.75ms
step:25/1555 train_time:865ms step_avg:34.59ms
step:26/1555 train_time:903ms step_avg:34.71ms
step:27/1555 train_time:934ms step_avg:34.58ms
step:28/1555 train_time:971ms step_avg:34.69ms
step:29/1555 train_time:1002ms step_avg:34.57ms
step:30/1555 train_time:1041ms step_avg:34.69ms
step:31/1555 train_time:1072ms step_avg:34.57ms
step:32/1555 train_time:1110ms step_avg:34.68ms
step:33/1555 train_time:1141ms step_avg:34.57ms
step:34/1555 train_time:1179ms step_avg:34.67ms
step:35/1555 train_time:1210ms step_avg:34.57ms
step:36/1555 train_time:1248ms step_avg:34.65ms
step:37/1555 train_time:1279ms step_avg:34.57ms
step:38/1555 train_time:1317ms step_avg:34.67ms
step:39/1555 train_time:1348ms step_avg:34.56ms
step:40/1555 train_time:1386ms step_avg:34.64ms
step:41/1555 train_time:1417ms step_avg:34.55ms
step:42/1555 train_time:1454ms step_avg:34.62ms
step:43/1555 train_time:1485ms step_avg:34.54ms
step:44/1555 train_time:1523ms step_avg:34.62ms
step:45/1555 train_time:1554ms step_avg:34.54ms
step:46/1555 train_time:1592ms step_avg:34.61ms
step:47/1555 train_time:1623ms step_avg:34.53ms
step:48/1555 train_time:1660ms step_avg:34.59ms
step:49/1555 train_time:1692ms step_avg:34.52ms
step:50/1555 train_time:1730ms step_avg:34.61ms
step:51/1555 train_time:1763ms step_avg:34.56ms
step:52/1555 train_time:1800ms step_avg:34.61ms
step:53/1555 train_time:1831ms step_avg:34.54ms
step:54/1555 train_time:1868ms step_avg:34.60ms
step:55/1555 train_time:1899ms step_avg:34.53ms
step:56/1555 train_time:1937ms step_avg:34.59ms
step:57/1555 train_time:1968ms step_avg:34.53ms
step:58/1555 train_time:2006ms step_avg:34.58ms
step:59/1555 train_time:2036ms step_avg:34.51ms
step:60/1555 train_time:2074ms step_avg:34.57ms
step:61/1555 train_time:2105ms step_avg:34.51ms
step:62/1555 train_time:2143ms step_avg:34.56ms
step:63/1555 train_time:2173ms step_avg:34.50ms
step:64/1555 train_time:2211ms step_avg:34.55ms
step:65/1555 train_time:2242ms step_avg:34.49ms
step:66/1555 train_time:2279ms step_avg:34.54ms
step:67/1555 train_time:2310ms step_avg:34.48ms
step:68/1555 train_time:2348ms step_avg:34.53ms
step:69/1555 train_time:2379ms step_avg:34.48ms
step:70/1555 train_time:2417ms step_avg:34.53ms
step:71/1555 train_time:2448ms step_avg:34.48ms
step:72/1555 train_time:2485ms step_avg:34.52ms
step:73/1555 train_time:2516ms step_avg:34.47ms
step:74/1555 train_time:2554ms step_avg:34.52ms
step:75/1555 train_time:2586ms step_avg:34.48ms
step:76/1555 train_time:2623ms step_avg:34.52ms
step:77/1555 train_time:2655ms step_avg:34.47ms
step:78/1555 train_time:2692ms step_avg:34.51ms
step:79/1555 train_time:2723ms step_avg:34.47ms
step:80/1555 train_time:2761ms step_avg:34.51ms
step:81/1555 train_time:2792ms step_avg:34.47ms
step:82/1555 train_time:2829ms step_avg:34.51ms
step:83/1555 train_time:2861ms step_avg:34.47ms
step:84/1555 train_time:2899ms step_avg:34.52ms
step:85/1555 train_time:2930ms step_avg:34.47ms
step:86/1555 train_time:2967ms step_avg:34.50ms
step:87/1555 train_time:2999ms step_avg:34.47ms
step:88/1555 train_time:3037ms step_avg:34.51ms
step:89/1555 train_time:3068ms step_avg:34.47ms
step:90/1555 train_time:3106ms step_avg:34.51ms
step:91/1555 train_time:3138ms step_avg:34.48ms
step:92/1555 train_time:3176ms step_avg:34.52ms
step:93/1555 train_time:3207ms step_avg:34.49ms
step:94/1555 train_time:3245ms step_avg:34.52ms
step:95/1555 train_time:3276ms step_avg:34.48ms
step:96/1555 train_time:3313ms step_avg:34.51ms
step:97/1555 train_time:3344ms step_avg:34.48ms
step:98/1555 train_time:3382ms step_avg:34.51ms
step:99/1555 train_time:3413ms step_avg:34.47ms
step:100/1555 train_time:3450ms step_avg:34.50ms
step:101/1555 train_time:3481ms step_avg:34.46ms
step:102/1555 train_time:3518ms step_avg:34.49ms
step:103/1555 train_time:3550ms step_avg:34.46ms
step:104/1555 train_time:3587ms step_avg:34.49ms
step:105/1555 train_time:3618ms step_avg:34.46ms
step:106/1555 train_time:3656ms step_avg:34.49ms
step:107/1555 train_time:3687ms step_avg:34.46ms
step:108/1555 train_time:3724ms step_avg:34.48ms
step:109/1555 train_time:3755ms step_avg:34.45ms
step:110/1555 train_time:3792ms step_avg:34.47ms
step:111/1555 train_time:3823ms step_avg:34.44ms
step:112/1555 train_time:3861ms step_avg:34.47ms
step:113/1555 train_time:3892ms step_avg:34.44ms
step:114/1555 train_time:3929ms step_avg:34.47ms
step:115/1555 train_time:3960ms step_avg:34.44ms
step:116/1555 train_time:3998ms step_avg:34.47ms
step:117/1555 train_time:4030ms step_avg:34.44ms
step:118/1555 train_time:4067ms step_avg:34.47ms
step:119/1555 train_time:4099ms step_avg:34.44ms
step:120/1555 train_time:4137ms step_avg:34.47ms
step:121/1555 train_time:4168ms step_avg:34.44ms
step:122/1555 train_time:4206ms step_avg:34.47ms
step:123/1555 train_time:4237ms step_avg:34.44ms
step:124/1555 train_time:4274ms step_avg:34.47ms
step:125/1555 train_time:4305ms step_avg:34.44ms
step:126/1555 train_time:4343ms step_avg:34.47ms
step:127/1555 train_time:4374ms step_avg:34.44ms
step:128/1555 train_time:4411ms step_avg:34.46ms
step:129/1555 train_time:4443ms step_avg:34.44ms
step:130/1555 train_time:4481ms step_avg:34.47ms
step:131/1555 train_time:4512ms step_avg:34.44ms
step:132/1555 train_time:4549ms step_avg:34.46ms
step:133/1555 train_time:4580ms step_avg:34.44ms
step:134/1555 train_time:4618ms step_avg:34.46ms
step:135/1555 train_time:4648ms step_avg:34.43ms
step:136/1555 train_time:4686ms step_avg:34.45ms
step:137/1555 train_time:4717ms step_avg:34.43ms
step:138/1555 train_time:4754ms step_avg:34.45ms
step:139/1555 train_time:4785ms step_avg:34.43ms
step:140/1555 train_time:4823ms step_avg:34.45ms
step:141/1555 train_time:4854ms step_avg:34.43ms
step:142/1555 train_time:4892ms step_avg:34.45ms
step:143/1555 train_time:4922ms step_avg:34.42ms
step:144/1555 train_time:4961ms step_avg:34.45ms
step:145/1555 train_time:4992ms step_avg:34.43ms
step:146/1555 train_time:5029ms step_avg:34.45ms
step:147/1555 train_time:5061ms step_avg:34.43ms
step:148/1555 train_time:5098ms step_avg:34.45ms
step:149/1555 train_time:5129ms step_avg:34.42ms
step:150/1555 train_time:5166ms step_avg:34.44ms
step:151/1555 train_time:5198ms step_avg:34.42ms
step:152/1555 train_time:5235ms step_avg:34.44ms
step:153/1555 train_time:5266ms step_avg:34.42ms
step:154/1555 train_time:5303ms step_avg:34.44ms
step:155/1555 train_time:5334ms step_avg:34.42ms
step:156/1555 train_time:5372ms step_avg:34.43ms
step:157/1555 train_time:5402ms step_avg:34.41ms
step:158/1555 train_time:5440ms step_avg:34.43ms
step:159/1555 train_time:5471ms step_avg:34.41ms
step:160/1555 train_time:5508ms step_avg:34.42ms
step:161/1555 train_time:5539ms step_avg:34.40ms
step:162/1555 train_time:5577ms step_avg:34.42ms
step:163/1555 train_time:5607ms step_avg:34.40ms
step:164/1555 train_time:5645ms step_avg:34.42ms
step:165/1555 train_time:5676ms step_avg:34.40ms
step:166/1555 train_time:5713ms step_avg:34.41ms
step:167/1555 train_time:5744ms step_avg:34.40ms
step:168/1555 train_time:5782ms step_avg:34.42ms
step:169/1555 train_time:5814ms step_avg:34.40ms
step:170/1555 train_time:5851ms step_avg:34.42ms
step:171/1555 train_time:5882ms step_avg:34.40ms
step:172/1555 train_time:5919ms step_avg:34.41ms
step:173/1555 train_time:5950ms step_avg:34.39ms
step:174/1555 train_time:5987ms step_avg:34.41ms
step:175/1555 train_time:6018ms step_avg:34.39ms
step:176/1555 train_time:6056ms step_avg:34.41ms
step:177/1555 train_time:6087ms step_avg:34.39ms
step:178/1555 train_time:6124ms step_avg:34.40ms
step:179/1555 train_time:6155ms step_avg:34.39ms
step:180/1555 train_time:6192ms step_avg:34.40ms
step:181/1555 train_time:6224ms step_avg:34.39ms
step:182/1555 train_time:6262ms step_avg:34.40ms
step:183/1555 train_time:6293ms step_avg:34.39ms
step:184/1555 train_time:6330ms step_avg:34.40ms
step:185/1555 train_time:6361ms step_avg:34.38ms
step:186/1555 train_time:6399ms step_avg:34.40ms
step:187/1555 train_time:6429ms step_avg:34.38ms
step:188/1555 train_time:6466ms step_avg:34.40ms
step:189/1555 train_time:6498ms step_avg:34.38ms
step:190/1555 train_time:6535ms step_avg:34.40ms
step:191/1555 train_time:6566ms step_avg:34.38ms
step:192/1555 train_time:6603ms step_avg:34.39ms
step:193/1555 train_time:6634ms step_avg:34.38ms
step:194/1555 train_time:6672ms step_avg:34.39ms
step:195/1555 train_time:6703ms step_avg:34.37ms
step:196/1555 train_time:6741ms step_avg:34.39ms
step:197/1555 train_time:6772ms step_avg:34.37ms
step:198/1555 train_time:6809ms step_avg:34.39ms
step:199/1555 train_time:6839ms step_avg:34.37ms
step:200/1555 train_time:6877ms step_avg:34.38ms
step:201/1555 train_time:6908ms step_avg:34.37ms
step:202/1555 train_time:6945ms step_avg:34.38ms
step:203/1555 train_time:6976ms step_avg:34.37ms
step:204/1555 train_time:7014ms step_avg:34.38ms
step:205/1555 train_time:7045ms step_avg:34.36ms
step:206/1555 train_time:7083ms step_avg:34.38ms
step:207/1555 train_time:7114ms step_avg:34.37ms
step:208/1555 train_time:7151ms step_avg:34.38ms
step:209/1555 train_time:7182ms step_avg:34.36ms
step:210/1555 train_time:7219ms step_avg:34.38ms
step:211/1555 train_time:7251ms step_avg:34.36ms
step:212/1555 train_time:7288ms step_avg:34.38ms
step:213/1555 train_time:7319ms step_avg:34.36ms
step:214/1555 train_time:7357ms step_avg:34.38ms
step:215/1555 train_time:7388ms step_avg:34.36ms
step:216/1555 train_time:7426ms step_avg:34.38ms
step:217/1555 train_time:7457ms step_avg:34.36ms
step:218/1555 train_time:7495ms step_avg:34.38ms
step:219/1555 train_time:7525ms step_avg:34.36ms
step:220/1555 train_time:7563ms step_avg:34.38ms
step:221/1555 train_time:7594ms step_avg:34.36ms
step:222/1555 train_time:7631ms step_avg:34.37ms
step:223/1555 train_time:7663ms step_avg:34.36ms
step:224/1555 train_time:7700ms step_avg:34.38ms
step:225/1555 train_time:7731ms step_avg:34.36ms
step:226/1555 train_time:7768ms step_avg:34.37ms
step:227/1555 train_time:7799ms step_avg:34.36ms
step:228/1555 train_time:7837ms step_avg:34.37ms
step:229/1555 train_time:7868ms step_avg:34.36ms
step:230/1555 train_time:7905ms step_avg:34.37ms
step:231/1555 train_time:7936ms step_avg:34.36ms
step:232/1555 train_time:7973ms step_avg:34.37ms
step:233/1555 train_time:8005ms step_avg:34.35ms
step:234/1555 train_time:8042ms step_avg:34.37ms
step:235/1555 train_time:8073ms step_avg:34.35ms
step:236/1555 train_time:8110ms step_avg:34.37ms
step:237/1555 train_time:8141ms step_avg:34.35ms
step:238/1555 train_time:8179ms step_avg:34.36ms
step:239/1555 train_time:8210ms step_avg:34.35ms
step:240/1555 train_time:8247ms step_avg:34.36ms
step:241/1555 train_time:8278ms step_avg:34.35ms
step:242/1555 train_time:8316ms step_avg:34.36ms
step:243/1555 train_time:8347ms step_avg:34.35ms
step:244/1555 train_time:8384ms step_avg:34.36ms
step:245/1555 train_time:8416ms step_avg:34.35ms
step:246/1555 train_time:8453ms step_avg:34.36ms
step:247/1555 train_time:8484ms step_avg:34.35ms
step:248/1555 train_time:8521ms step_avg:34.36ms
step:249/1555 train_time:8553ms step_avg:34.35ms
step:250/1555 train_time:8590ms step_avg:34.36ms
step:250/1555 val_loss:4.5471 train_time:8640ms step_avg:34.56ms
step:251/1555 train_time:8658ms step_avg:34.50ms
step:252/1555 train_time:8678ms step_avg:34.44ms
step:253/1555 train_time:8695ms step_avg:34.37ms
step:254/1555 train_time:8732ms step_avg:34.38ms
step:255/1555 train_time:8765ms step_avg:34.37ms
step:256/1555 train_time:8804ms step_avg:34.39ms
step:257/1555 train_time:8835ms step_avg:34.38ms
step:258/1555 train_time:8873ms step_avg:34.39ms
step:259/1555 train_time:8904ms step_avg:34.38ms
step:260/1555 train_time:8941ms step_avg:34.39ms
step:261/1555 train_time:8972ms step_avg:34.38ms
step:262/1555 train_time:9010ms step_avg:34.39ms
step:263/1555 train_time:9041ms step_avg:34.38ms
step:264/1555 train_time:9079ms step_avg:34.39ms
step:265/1555 train_time:9109ms step_avg:34.37ms
step:266/1555 train_time:9146ms step_avg:34.38ms
step:267/1555 train_time:9177ms step_avg:34.37ms
step:268/1555 train_time:9215ms step_avg:34.38ms
step:269/1555 train_time:9246ms step_avg:34.37ms
step:270/1555 train_time:9283ms step_avg:34.38ms
step:271/1555 train_time:9314ms step_avg:34.37ms
step:272/1555 train_time:9351ms step_avg:34.38ms
step:273/1555 train_time:9382ms step_avg:34.37ms
step:274/1555 train_time:9419ms step_avg:34.38ms
step:275/1555 train_time:9450ms step_avg:34.36ms
step:276/1555 train_time:9487ms step_avg:34.37ms
step:277/1555 train_time:9518ms step_avg:34.36ms
step:278/1555 train_time:9555ms step_avg:34.37ms
step:279/1555 train_time:9586ms step_avg:34.36ms
step:280/1555 train_time:9623ms step_avg:34.37ms
step:281/1555 train_time:9654ms step_avg:34.36ms
step:282/1555 train_time:9693ms step_avg:34.37ms
step:283/1555 train_time:9724ms step_avg:34.36ms
step:284/1555 train_time:9762ms step_avg:34.37ms
step:285/1555 train_time:9793ms step_avg:34.36ms
step:286/1555 train_time:9831ms step_avg:34.37ms
step:287/1555 train_time:9862ms step_avg:34.36ms
step:288/1555 train_time:9899ms step_avg:34.37ms
step:289/1555 train_time:9931ms step_avg:34.36ms
step:290/1555 train_time:9968ms step_avg:34.37ms
step:291/1555 train_time:9999ms step_avg:34.36ms
step:292/1555 train_time:10037ms step_avg:34.37ms
step:293/1555 train_time:10068ms step_avg:34.36ms
step:294/1555 train_time:10106ms step_avg:34.37ms
step:295/1555 train_time:10137ms step_avg:34.36ms
step:296/1555 train_time:10175ms step_avg:34.37ms
step:297/1555 train_time:10206ms step_avg:34.36ms
step:298/1555 train_time:10242ms step_avg:34.37ms
step:299/1555 train_time:10274ms step_avg:34.36ms
step:300/1555 train_time:10311ms step_avg:34.37ms
step:301/1555 train_time:10342ms step_avg:34.36ms
step:302/1555 train_time:10380ms step_avg:34.37ms
step:303/1555 train_time:10410ms step_avg:34.36ms
step:304/1555 train_time:10448ms step_avg:34.37ms
step:305/1555 train_time:10479ms step_avg:34.36ms
step:306/1555 train_time:10516ms step_avg:34.37ms
step:307/1555 train_time:10547ms step_avg:34.36ms
step:308/1555 train_time:10584ms step_avg:34.36ms
step:309/1555 train_time:10615ms step_avg:34.35ms
step:310/1555 train_time:10653ms step_avg:34.36ms
step:311/1555 train_time:10683ms step_avg:34.35ms
step:312/1555 train_time:10721ms step_avg:34.36ms
step:313/1555 train_time:10752ms step_avg:34.35ms
step:314/1555 train_time:10789ms step_avg:34.36ms
step:315/1555 train_time:10820ms step_avg:34.35ms
step:316/1555 train_time:10858ms step_avg:34.36ms
step:317/1555 train_time:10889ms step_avg:34.35ms
step:318/1555 train_time:10926ms step_avg:34.36ms
step:319/1555 train_time:10958ms step_avg:34.35ms
step:320/1555 train_time:10995ms step_avg:34.36ms
step:321/1555 train_time:11026ms step_avg:34.35ms
step:322/1555 train_time:11063ms step_avg:34.36ms
step:323/1555 train_time:11094ms step_avg:34.35ms
step:324/1555 train_time:11132ms step_avg:34.36ms
step:325/1555 train_time:11163ms step_avg:34.35ms
step:326/1555 train_time:11200ms step_avg:34.36ms
step:327/1555 train_time:11231ms step_avg:34.35ms
step:328/1555 train_time:11269ms step_avg:34.36ms
step:329/1555 train_time:11300ms step_avg:34.35ms
step:330/1555 train_time:11337ms step_avg:34.35ms
step:331/1555 train_time:11367ms step_avg:34.34ms
step:332/1555 train_time:11405ms step_avg:34.35ms
step:333/1555 train_time:11435ms step_avg:34.34ms
step:334/1555 train_time:11473ms step_avg:34.35ms
step:335/1555 train_time:11504ms step_avg:34.34ms
step:336/1555 train_time:11541ms step_avg:34.35ms
step:337/1555 train_time:11572ms step_avg:34.34ms
step:338/1555 train_time:11610ms step_avg:34.35ms
step:339/1555 train_time:11641ms step_avg:34.34ms
step:340/1555 train_time:11679ms step_avg:34.35ms
step:341/1555 train_time:11710ms step_avg:34.34ms
step:342/1555 train_time:11748ms step_avg:34.35ms
step:343/1555 train_time:11779ms step_avg:34.34ms
step:344/1555 train_time:11817ms step_avg:34.35ms
step:345/1555 train_time:11848ms step_avg:34.34ms
step:346/1555 train_time:11885ms step_avg:34.35ms
step:347/1555 train_time:11916ms step_avg:34.34ms
step:348/1555 train_time:11954ms step_avg:34.35ms
step:349/1555 train_time:11985ms step_avg:34.34ms
step:350/1555 train_time:12022ms step_avg:34.35ms
step:351/1555 train_time:12053ms step_avg:34.34ms
step:352/1555 train_time:12090ms step_avg:34.35ms
step:353/1555 train_time:12122ms step_avg:34.34ms
step:354/1555 train_time:12159ms step_avg:34.35ms
step:355/1555 train_time:12189ms step_avg:34.34ms
step:356/1555 train_time:12227ms step_avg:34.34ms
step:357/1555 train_time:12258ms step_avg:34.34ms
step:358/1555 train_time:12296ms step_avg:34.35ms
step:359/1555 train_time:12326ms step_avg:34.34ms
step:360/1555 train_time:12364ms step_avg:34.34ms
step:361/1555 train_time:12394ms step_avg:34.33ms
step:362/1555 train_time:12432ms step_avg:34.34ms
step:363/1555 train_time:12463ms step_avg:34.33ms
step:364/1555 train_time:12500ms step_avg:34.34ms
step:365/1555 train_time:12532ms step_avg:34.33ms
step:366/1555 train_time:12569ms step_avg:34.34ms
step:367/1555 train_time:12600ms step_avg:34.33ms
step:368/1555 train_time:12637ms step_avg:34.34ms
step:369/1555 train_time:12668ms step_avg:34.33ms
step:370/1555 train_time:12705ms step_avg:34.34ms
step:371/1555 train_time:12737ms step_avg:34.33ms
step:372/1555 train_time:12774ms step_avg:34.34ms
step:373/1555 train_time:12805ms step_avg:34.33ms
step:374/1555 train_time:12842ms step_avg:34.34ms
step:375/1555 train_time:12873ms step_avg:34.33ms
step:376/1555 train_time:12910ms step_avg:34.34ms
step:377/1555 train_time:12942ms step_avg:34.33ms
step:378/1555 train_time:12979ms step_avg:34.34ms
step:379/1555 train_time:13010ms step_avg:34.33ms
step:380/1555 train_time:13048ms step_avg:34.34ms
step:381/1555 train_time:13079ms step_avg:34.33ms
step:382/1555 train_time:13116ms step_avg:34.34ms
step:383/1555 train_time:13147ms step_avg:34.33ms
step:384/1555 train_time:13184ms step_avg:34.33ms
step:385/1555 train_time:13216ms step_avg:34.33ms
step:386/1555 train_time:13253ms step_avg:34.33ms
step:387/1555 train_time:13284ms step_avg:34.33ms
step:388/1555 train_time:13322ms step_avg:34.33ms
step:389/1555 train_time:13353ms step_avg:34.33ms
step:390/1555 train_time:13391ms step_avg:34.34ms
step:391/1555 train_time:13422ms step_avg:34.33ms
step:392/1555 train_time:13459ms step_avg:34.33ms
step:393/1555 train_time:13490ms step_avg:34.33ms
step:394/1555 train_time:13527ms step_avg:34.33ms
step:395/1555 train_time:13558ms step_avg:34.32ms
step:396/1555 train_time:13596ms step_avg:34.33ms
step:397/1555 train_time:13626ms step_avg:34.32ms
step:398/1555 train_time:13664ms step_avg:34.33ms
step:399/1555 train_time:13695ms step_avg:34.32ms
step:400/1555 train_time:13732ms step_avg:34.33ms
step:401/1555 train_time:13763ms step_avg:34.32ms
step:402/1555 train_time:13800ms step_avg:34.33ms
step:403/1555 train_time:13831ms step_avg:34.32ms
step:404/1555 train_time:13868ms step_avg:34.33ms
step:405/1555 train_time:13899ms step_avg:34.32ms
step:406/1555 train_time:13936ms step_avg:34.33ms
step:407/1555 train_time:13967ms step_avg:34.32ms
step:408/1555 train_time:14004ms step_avg:34.32ms
step:409/1555 train_time:14035ms step_avg:34.32ms
step:410/1555 train_time:14073ms step_avg:34.32ms
step:411/1555 train_time:14104ms step_avg:34.32ms
step:412/1555 train_time:14141ms step_avg:34.32ms
step:413/1555 train_time:14172ms step_avg:34.32ms
step:414/1555 train_time:14209ms step_avg:34.32ms
step:415/1555 train_time:14241ms step_avg:34.31ms
step:416/1555 train_time:14278ms step_avg:34.32ms
step:417/1555 train_time:14309ms step_avg:34.31ms
step:418/1555 train_time:14346ms step_avg:34.32ms
step:419/1555 train_time:14378ms step_avg:34.31ms
step:420/1555 train_time:14415ms step_avg:34.32ms
step:421/1555 train_time:14446ms step_avg:34.31ms
step:422/1555 train_time:14483ms step_avg:34.32ms
step:423/1555 train_time:14514ms step_avg:34.31ms
step:424/1555 train_time:14552ms step_avg:34.32ms
step:425/1555 train_time:14583ms step_avg:34.31ms
step:426/1555 train_time:14620ms step_avg:34.32ms
step:427/1555 train_time:14651ms step_avg:34.31ms
step:428/1555 train_time:14689ms step_avg:34.32ms
step:429/1555 train_time:14719ms step_avg:34.31ms
step:430/1555 train_time:14756ms step_avg:34.32ms
step:431/1555 train_time:14787ms step_avg:34.31ms
step:432/1555 train_time:14824ms step_avg:34.32ms
step:433/1555 train_time:14855ms step_avg:34.31ms
step:434/1555 train_time:14893ms step_avg:34.31ms
step:435/1555 train_time:14924ms step_avg:34.31ms
step:436/1555 train_time:14961ms step_avg:34.31ms
step:437/1555 train_time:14993ms step_avg:34.31ms
step:438/1555 train_time:15030ms step_avg:34.32ms
step:439/1555 train_time:15061ms step_avg:34.31ms
step:440/1555 train_time:15098ms step_avg:34.31ms
step:441/1555 train_time:15129ms step_avg:34.31ms
step:442/1555 train_time:15167ms step_avg:34.32ms
step:443/1555 train_time:15199ms step_avg:34.31ms
step:444/1555 train_time:15236ms step_avg:34.32ms
step:445/1555 train_time:15267ms step_avg:34.31ms
step:446/1555 train_time:15304ms step_avg:34.31ms
step:447/1555 train_time:15335ms step_avg:34.31ms
step:448/1555 train_time:15372ms step_avg:34.31ms
step:449/1555 train_time:15403ms step_avg:34.31ms
step:450/1555 train_time:15441ms step_avg:34.31ms
step:451/1555 train_time:15472ms step_avg:34.31ms
step:452/1555 train_time:15509ms step_avg:34.31ms
step:453/1555 train_time:15540ms step_avg:34.30ms
step:454/1555 train_time:15577ms step_avg:34.31ms
step:455/1555 train_time:15609ms step_avg:34.30ms
step:456/1555 train_time:15646ms step_avg:34.31ms
step:457/1555 train_time:15677ms step_avg:34.30ms
step:458/1555 train_time:15714ms step_avg:34.31ms
step:459/1555 train_time:15745ms step_avg:34.30ms
step:460/1555 train_time:15782ms step_avg:34.31ms
step:461/1555 train_time:15814ms step_avg:34.30ms
step:462/1555 train_time:15852ms step_avg:34.31ms
step:463/1555 train_time:15882ms step_avg:34.30ms
step:464/1555 train_time:15920ms step_avg:34.31ms
step:465/1555 train_time:15951ms step_avg:34.30ms
step:466/1555 train_time:15989ms step_avg:34.31ms
step:467/1555 train_time:16021ms step_avg:34.31ms
step:468/1555 train_time:16058ms step_avg:34.31ms
step:469/1555 train_time:16089ms step_avg:34.31ms
step:470/1555 train_time:16127ms step_avg:34.31ms
step:471/1555 train_time:16158ms step_avg:34.31ms
step:472/1555 train_time:16195ms step_avg:34.31ms
step:473/1555 train_time:16226ms step_avg:34.30ms
step:474/1555 train_time:16263ms step_avg:34.31ms
step:475/1555 train_time:16294ms step_avg:34.30ms
step:476/1555 train_time:16332ms step_avg:34.31ms
step:477/1555 train_time:16363ms step_avg:34.30ms
step:478/1555 train_time:16400ms step_avg:34.31ms
step:479/1555 train_time:16431ms step_avg:34.30ms
step:480/1555 train_time:16469ms step_avg:34.31ms
step:481/1555 train_time:16500ms step_avg:34.30ms
step:482/1555 train_time:16537ms step_avg:34.31ms
step:483/1555 train_time:16568ms step_avg:34.30ms
step:484/1555 train_time:16605ms step_avg:34.31ms
step:485/1555 train_time:16636ms step_avg:34.30ms
step:486/1555 train_time:16673ms step_avg:34.31ms
step:487/1555 train_time:16704ms step_avg:34.30ms
step:488/1555 train_time:16741ms step_avg:34.31ms
step:489/1555 train_time:16772ms step_avg:34.30ms
step:490/1555 train_time:16810ms step_avg:34.31ms
step:491/1555 train_time:16841ms step_avg:34.30ms
step:492/1555 train_time:16878ms step_avg:34.31ms
step:493/1555 train_time:16909ms step_avg:34.30ms
step:494/1555 train_time:16947ms step_avg:34.31ms
step:495/1555 train_time:16978ms step_avg:34.30ms
step:496/1555 train_time:17016ms step_avg:34.31ms
step:497/1555 train_time:17047ms step_avg:34.30ms
step:498/1555 train_time:17084ms step_avg:34.30ms
step:499/1555 train_time:17115ms step_avg:34.30ms
step:500/1555 train_time:17153ms step_avg:34.31ms
step:500/1555 val_loss:4.2237 train_time:17203ms step_avg:34.41ms
step:501/1555 train_time:17221ms step_avg:34.37ms
step:502/1555 train_time:17240ms step_avg:34.34ms
step:503/1555 train_time:17256ms step_avg:34.31ms
step:504/1555 train_time:17292ms step_avg:34.31ms
step:505/1555 train_time:17326ms step_avg:34.31ms
step:506/1555 train_time:17368ms step_avg:34.32ms
step:507/1555 train_time:17423ms step_avg:34.36ms
step:508/1555 train_time:17487ms step_avg:34.42ms
step:509/1555 train_time:17544ms step_avg:34.47ms
step:510/1555 train_time:17607ms step_avg:34.52ms
step:511/1555 train_time:17665ms step_avg:34.57ms
step:512/1555 train_time:17729ms step_avg:34.63ms
step:513/1555 train_time:17785ms step_avg:34.67ms
step:514/1555 train_time:17848ms step_avg:34.72ms
step:515/1555 train_time:17905ms step_avg:34.77ms
step:516/1555 train_time:17968ms step_avg:34.82ms
step:517/1555 train_time:18025ms step_avg:34.86ms
step:518/1555 train_time:18089ms step_avg:34.92ms
step:519/1555 train_time:18146ms step_avg:34.96ms
step:520/1555 train_time:18211ms step_avg:35.02ms
step:521/1555 train_time:18270ms step_avg:35.07ms
step:522/1555 train_time:18336ms step_avg:35.13ms
step:523/1555 train_time:18395ms step_avg:35.17ms
step:524/1555 train_time:18461ms step_avg:35.23ms
step:525/1555 train_time:18519ms step_avg:35.27ms
step:526/1555 train_time:18583ms step_avg:35.33ms
step:527/1555 train_time:18641ms step_avg:35.37ms
step:528/1555 train_time:18704ms step_avg:35.42ms
step:529/1555 train_time:18761ms step_avg:35.47ms
step:530/1555 train_time:18825ms step_avg:35.52ms
step:531/1555 train_time:18883ms step_avg:35.56ms
step:532/1555 train_time:18946ms step_avg:35.61ms
step:533/1555 train_time:19003ms step_avg:35.65ms
step:534/1555 train_time:19067ms step_avg:35.71ms
step:535/1555 train_time:19125ms step_avg:35.75ms
step:536/1555 train_time:19191ms step_avg:35.80ms
step:537/1555 train_time:19250ms step_avg:35.85ms
step:538/1555 train_time:19315ms step_avg:35.90ms
step:539/1555 train_time:19372ms step_avg:35.94ms
step:540/1555 train_time:19437ms step_avg:35.99ms
step:541/1555 train_time:19495ms step_avg:36.04ms
step:542/1555 train_time:19560ms step_avg:36.09ms
step:543/1555 train_time:19618ms 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:19927ms step_avg:36.36ms
step:549/1555 train_time:19984ms step_avg:36.40ms
step:550/1555 train_time:20049ms step_avg:36.45ms
step:551/1555 train_time:20107ms step_avg:36.49ms
step:552/1555 train_time:20170ms step_avg:36.54ms
step:553/1555 train_time:20229ms step_avg:36.58ms
step:554/1555 train_time:20292ms step_avg:36.63ms
step:555/1555 train_time:20350ms step_avg:36.67ms
step:556/1555 train_time:20415ms step_avg:36.72ms
step:557/1555 train_time:20472ms step_avg:36.75ms
step:558/1555 train_time:20538ms step_avg:36.81ms
step:559/1555 train_time:20597ms step_avg:36.85ms
step:560/1555 train_time:20661ms step_avg:36.89ms
step:561/1555 train_time:20719ms step_avg:36.93ms
step:562/1555 train_time:20783ms step_avg:36.98ms
step:563/1555 train_time:20842ms step_avg:37.02ms
step:564/1555 train_time:20905ms step_avg:37.07ms
step:565/1555 train_time:20963ms step_avg:37.10ms
step:566/1555 train_time:21028ms step_avg:37.15ms
step:567/1555 train_time:21086ms step_avg:37.19ms
step:568/1555 train_time:21150ms step_avg:37.24ms
step:569/1555 train_time:21208ms step_avg:37.27ms
step:570/1555 train_time:21271ms step_avg:37.32ms
step:571/1555 train_time:21329ms step_avg:37.35ms
step:572/1555 train_time:21393ms step_avg:37.40ms
step:573/1555 train_time:21451ms step_avg:37.44ms
step:574/1555 train_time:21516ms step_avg:37.48ms
step:575/1555 train_time:21574ms step_avg:37.52ms
step:576/1555 train_time:21639ms step_avg:37.57ms
step:577/1555 train_time:21697ms step_avg:37.60ms
step:578/1555 train_time:21762ms step_avg:37.65ms
step:579/1555 train_time:21820ms step_avg:37.69ms
step:580/1555 train_time:21884ms step_avg:37.73ms
step:581/1555 train_time:21942ms step_avg:37.77ms
step:582/1555 train_time:22007ms step_avg:37.81ms
step:583/1555 train_time:22065ms step_avg:37.85ms
step:584/1555 train_time:22129ms step_avg:37.89ms
step:585/1555 train_time:22187ms step_avg:37.93ms
step:586/1555 train_time:22251ms step_avg:37.97ms
step:587/1555 train_time:22308ms step_avg:38.00ms
step:588/1555 train_time:22373ms step_avg:38.05ms
step:589/1555 train_time:22431ms step_avg:38.08ms
step:590/1555 train_time:22494ms step_avg:38.13ms
step:591/1555 train_time:22552ms step_avg:38.16ms
step:592/1555 train_time:22618ms step_avg:38.21ms
step:593/1555 train_time:22676ms step_avg:38.24ms
step:594/1555 train_time:22741ms step_avg:38.28ms
step:595/1555 train_time:22799ms step_avg:38.32ms
step:596/1555 train_time:22863ms step_avg:38.36ms
step:597/1555 train_time:22920ms step_avg:38.39ms
step:598/1555 train_time:22985ms step_avg:38.44ms
step:599/1555 train_time:23044ms step_avg:38.47ms
step:600/1555 train_time:23108ms step_avg:38.51ms
step:601/1555 train_time:23166ms step_avg:38.55ms
step:602/1555 train_time:23230ms step_avg:38.59ms
step:603/1555 train_time:23288ms step_avg:38.62ms
step:604/1555 train_time:23352ms step_avg:38.66ms
step:605/1555 train_time:23410ms step_avg:38.69ms
step:606/1555 train_time:23473ms step_avg:38.73ms
step:607/1555 train_time:23530ms step_avg:38.77ms
step:608/1555 train_time:23595ms step_avg:38.81ms
step:609/1555 train_time:23652ms step_avg:38.84ms
step:610/1555 train_time:23718ms step_avg:38.88ms
step:611/1555 train_time:23775ms step_avg:38.91ms
step:612/1555 train_time:23840ms step_avg:38.95ms
step:613/1555 train_time:23898ms step_avg:38.99ms
step:614/1555 train_time:23963ms step_avg:39.03ms
step:615/1555 train_time:24021ms step_avg:39.06ms
step:616/1555 train_time:24086ms step_avg:39.10ms
step:617/1555 train_time:24145ms step_avg:39.13ms
step:618/1555 train_time:24209ms step_avg:39.17ms
step:619/1555 train_time:24266ms step_avg:39.20ms
step:620/1555 train_time:24330ms step_avg:39.24ms
step:621/1555 train_time:24388ms step_avg:39.27ms
step:622/1555 train_time:24452ms step_avg:39.31ms
step:623/1555 train_time:24509ms step_avg:39.34ms
step:624/1555 train_time:24573ms step_avg:39.38ms
step:625/1555 train_time:24630ms step_avg:39.41ms
step:626/1555 train_time:24696ms step_avg:39.45ms
step:627/1555 train_time:24754ms step_avg:39.48ms
step:628/1555 train_time:24819ms step_avg:39.52ms
step:629/1555 train_time:24877ms step_avg:39.55ms
step:630/1555 train_time:24941ms step_avg:39.59ms
step:631/1555 train_time:25000ms step_avg:39.62ms
step:632/1555 train_time:25065ms step_avg:39.66ms
step:633/1555 train_time:25123ms step_avg:39.69ms
step:634/1555 train_time:25188ms step_avg:39.73ms
step:635/1555 train_time:25246ms step_avg:39.76ms
step:636/1555 train_time:25309ms step_avg:39.79ms
step:637/1555 train_time:25367ms step_avg:39.82ms
step:638/1555 train_time:25430ms step_avg:39.86ms
step:639/1555 train_time:25488ms step_avg:39.89ms
step:640/1555 train_time:25553ms step_avg:39.93ms
step:641/1555 train_time:25610ms step_avg:39.95ms
step:642/1555 train_time:25674ms step_avg:39.99ms
step:643/1555 train_time:25732ms step_avg:40.02ms
step:644/1555 train_time:25796ms step_avg:40.06ms
step:645/1555 train_time:25855ms step_avg:40.09ms
step:646/1555 train_time:25921ms step_avg:40.12ms
step:647/1555 train_time:25979ms step_avg:40.15ms
step:648/1555 train_time:26044ms step_avg:40.19ms
step:649/1555 train_time:26101ms step_avg:40.22ms
step:650/1555 train_time:26165ms step_avg:40.25ms
step:651/1555 train_time:26223ms step_avg:40.28ms
step:652/1555 train_time:26287ms step_avg:40.32ms
step:653/1555 train_time:26347ms step_avg:40.35ms
step:654/1555 train_time:26410ms step_avg:40.38ms
step:655/1555 train_time:26468ms step_avg:40.41ms
step:656/1555 train_time:26532ms step_avg:40.44ms
step:657/1555 train_time:26589ms step_avg:40.47ms
step:658/1555 train_time:26654ms step_avg:40.51ms
step:659/1555 train_time:26711ms step_avg:40.53ms
step:660/1555 train_time:26776ms step_avg:40.57ms
step:661/1555 train_time:26834ms step_avg:40.60ms
step:662/1555 train_time:26899ms step_avg:40.63ms
step:663/1555 train_time:26957ms step_avg:40.66ms
step:664/1555 train_time:27022ms step_avg:40.70ms
step:665/1555 train_time:27079ms step_avg:40.72ms
step:666/1555 train_time:27144ms step_avg:40.76ms
step:667/1555 train_time:27201ms step_avg:40.78ms
step:668/1555 train_time:27266ms step_avg:40.82ms
step:669/1555 train_time:27324ms step_avg:40.84ms
step:670/1555 train_time:27388ms step_avg:40.88ms
step:671/1555 train_time:27446ms step_avg:40.90ms
step:672/1555 train_time:27509ms step_avg:40.94ms
step:673/1555 train_time:27567ms step_avg:40.96ms
step:674/1555 train_time:27631ms step_avg:41.00ms
step:675/1555 train_time:27688ms step_avg:41.02ms
step:676/1555 train_time:27753ms step_avg:41.05ms
step:677/1555 train_time:27810ms step_avg:41.08ms
step:678/1555 train_time:27875ms step_avg:41.11ms
step:679/1555 train_time:27933ms step_avg:41.14ms
step:680/1555 train_time:27998ms step_avg:41.17ms
step:681/1555 train_time:28056ms step_avg:41.20ms
step:682/1555 train_time:28121ms step_avg:41.23ms
step:683/1555 train_time:28178ms step_avg:41.26ms
step:684/1555 train_time:28242ms step_avg:41.29ms
step:685/1555 train_time:28300ms step_avg:41.31ms
step:686/1555 train_time:28365ms step_avg:41.35ms
step:687/1555 train_time:28423ms step_avg:41.37ms
step:688/1555 train_time:28487ms step_avg:41.41ms
step:689/1555 train_time:28546ms step_avg:41.43ms
step:690/1555 train_time:28609ms step_avg:41.46ms
step:691/1555 train_time:28667ms step_avg:41.49ms
step:692/1555 train_time:28731ms step_avg:41.52ms
step:693/1555 train_time:28788ms step_avg:41.54ms
step:694/1555 train_time:28853ms step_avg:41.57ms
step:695/1555 train_time:28909ms step_avg:41.60ms
step:696/1555 train_time:28975ms step_avg:41.63ms
step:697/1555 train_time:29033ms step_avg:41.65ms
step:698/1555 train_time:29098ms step_avg:41.69ms
step:699/1555 train_time:29155ms step_avg:41.71ms
step:700/1555 train_time:29220ms step_avg:41.74ms
step:701/1555 train_time:29277ms step_avg:41.77ms
step:702/1555 train_time:29342ms step_avg:41.80ms
step:703/1555 train_time:29400ms step_avg:41.82ms
step:704/1555 train_time:29465ms step_avg:41.85ms
step:705/1555 train_time:29522ms step_avg:41.88ms
step:706/1555 train_time:29586ms step_avg:41.91ms
step:707/1555 train_time:29644ms step_avg:41.93ms
step:708/1555 train_time:29710ms step_avg:41.96ms
step:709/1555 train_time:29767ms step_avg:41.98ms
step:710/1555 train_time:29831ms step_avg:42.02ms
step:711/1555 train_time:29888ms step_avg:42.04ms
step:712/1555 train_time:29952ms step_avg:42.07ms
step:713/1555 train_time:30010ms step_avg:42.09ms
step:714/1555 train_time:30075ms step_avg:42.12ms
step:715/1555 train_time:30132ms step_avg:42.14ms
step:716/1555 train_time:30198ms step_avg:42.18ms
step:717/1555 train_time:30256ms step_avg:42.20ms
step:718/1555 train_time:30319ms step_avg:42.23ms
step:719/1555 train_time:30377ms step_avg:42.25ms
step:720/1555 train_time:30442ms step_avg:42.28ms
step:721/1555 train_time:30500ms step_avg:42.30ms
step:722/1555 train_time:30564ms step_avg:42.33ms
step:723/1555 train_time:30623ms step_avg:42.35ms
step:724/1555 train_time:30687ms step_avg:42.39ms
step:725/1555 train_time:30745ms step_avg:42.41ms
step:726/1555 train_time:30810ms step_avg:42.44ms
step:727/1555 train_time:30867ms step_avg:42.46ms
step:728/1555 train_time:30931ms step_avg:42.49ms
step:729/1555 train_time:30988ms step_avg:42.51ms
step:730/1555 train_time:31054ms step_avg:42.54ms
step:731/1555 train_time:31110ms step_avg:42.56ms
step:732/1555 train_time:31176ms step_avg:42.59ms
step:733/1555 train_time:31233ms step_avg:42.61ms
step:734/1555 train_time:31298ms step_avg:42.64ms
step:735/1555 train_time:31356ms step_avg:42.66ms
step:736/1555 train_time:31420ms step_avg:42.69ms
step:737/1555 train_time:31478ms step_avg:42.71ms
step:738/1555 train_time:31543ms step_avg:42.74ms
step:739/1555 train_time:31602ms step_avg:42.76ms
step:740/1555 train_time:31665ms step_avg:42.79ms
step:741/1555 train_time:31724ms step_avg:42.81ms
step:742/1555 train_time:31788ms step_avg:42.84ms
step:743/1555 train_time:31846ms step_avg:42.86ms
step:744/1555 train_time:31910ms step_avg:42.89ms
step:745/1555 train_time:31967ms step_avg:42.91ms
step:746/1555 train_time:32032ms step_avg:42.94ms
step:747/1555 train_time:32089ms step_avg:42.96ms
step:748/1555 train_time:32154ms step_avg:42.99ms
step:749/1555 train_time:32211ms step_avg:43.00ms
step:750/1555 train_time:32276ms step_avg:43.03ms
step:750/1555 val_loss:3.8786 train_time:32358ms step_avg:43.14ms
step:751/1555 train_time:32377ms step_avg:43.11ms
step:752/1555 train_time:32399ms step_avg:43.08ms
step:753/1555 train_time:32458ms step_avg:43.11ms
step:754/1555 train_time:32527ms step_avg:43.14ms
step:755/1555 train_time:32588ms step_avg:43.16ms
step:756/1555 train_time:32652ms step_avg:43.19ms
step:757/1555 train_time:32709ms step_avg:43.21ms
step:758/1555 train_time:32773ms step_avg:43.24ms
step:759/1555 train_time:32830ms step_avg:43.25ms
step:760/1555 train_time:32894ms step_avg:43.28ms
step:761/1555 train_time:32950ms step_avg:43.30ms
step:762/1555 train_time:33016ms step_avg:43.33ms
step:763/1555 train_time:33073ms step_avg:43.35ms
step:764/1555 train_time:33137ms step_avg:43.37ms
step:765/1555 train_time:33194ms step_avg:43.39ms
step:766/1555 train_time:33257ms step_avg:43.42ms
step:767/1555 train_time:33315ms step_avg:43.44ms
step:768/1555 train_time:33379ms step_avg:43.46ms
step:769/1555 train_time:33439ms step_avg:43.48ms
step:770/1555 train_time:33504ms step_avg:43.51ms
step:771/1555 train_time:33564ms step_avg:43.53ms
step:772/1555 train_time:33629ms step_avg:43.56ms
step:773/1555 train_time:33686ms step_avg:43.58ms
step:774/1555 train_time:33750ms step_avg:43.60ms
step:775/1555 train_time:33807ms step_avg:43.62ms
step:776/1555 train_time:33870ms step_avg:43.65ms
step:777/1555 train_time:33927ms step_avg:43.66ms
step:778/1555 train_time:33991ms step_avg:43.69ms
step:779/1555 train_time:34049ms step_avg:43.71ms
step:780/1555 train_time:34112ms step_avg:43.73ms
step:781/1555 train_time:34170ms step_avg:43.75ms
step:782/1555 train_time:34233ms step_avg:43.78ms
step:783/1555 train_time:34290ms step_avg:43.79ms
step:784/1555 train_time:34355ms step_avg:43.82ms
step:785/1555 train_time:34413ms step_avg:43.84ms
step:786/1555 train_time:34479ms step_avg:43.87ms
step:787/1555 train_time:34537ms step_avg:43.88ms
step:788/1555 train_time:34602ms step_avg:43.91ms
step:789/1555 train_time:34660ms step_avg:43.93ms
step:790/1555 train_time:34724ms step_avg:43.96ms
step:791/1555 train_time:34782ms step_avg:43.97ms
step:792/1555 train_time:34848ms step_avg:44.00ms
step:793/1555 train_time:34906ms step_avg:44.02ms
step:794/1555 train_time:34971ms step_avg:44.04ms
step:795/1555 train_time:35028ms step_avg:44.06ms
step:796/1555 train_time:35091ms step_avg:44.08ms
step:797/1555 train_time:35149ms step_avg:44.10ms
step:798/1555 train_time:35212ms step_avg:44.13ms
step:799/1555 train_time:35270ms step_avg:44.14ms
step:800/1555 train_time:35333ms step_avg:44.17ms
step:801/1555 train_time:35390ms step_avg:44.18ms
step:802/1555 train_time:35455ms step_avg:44.21ms
step:803/1555 train_time:35513ms step_avg:44.23ms
step:804/1555 train_time:35579ms step_avg:44.25ms
step:805/1555 train_time:35637ms step_avg:44.27ms
step:806/1555 train_time:35701ms step_avg:44.29ms
step:807/1555 train_time:35760ms step_avg:44.31ms
step:808/1555 train_time:35823ms step_avg:44.34ms
step:809/1555 train_time:35882ms step_avg:44.35ms
step:810/1555 train_time:35947ms step_avg:44.38ms
step:811/1555 train_time:36005ms step_avg:44.40ms
step:812/1555 train_time:36069ms step_avg:44.42ms
step:813/1555 train_time:36126ms step_avg:44.44ms
step:814/1555 train_time:36189ms step_avg:44.46ms
step:815/1555 train_time:36248ms step_avg:44.48ms
step:816/1555 train_time:36312ms step_avg:44.50ms
step:817/1555 train_time:36370ms step_avg:44.52ms
step:818/1555 train_time:36433ms step_avg:44.54ms
step:819/1555 train_time:36491ms step_avg:44.56ms
step:820/1555 train_time:36556ms step_avg:44.58ms
step:821/1555 train_time:36613ms step_avg:44.60ms
step:822/1555 train_time:36679ms step_avg:44.62ms
step:823/1555 train_time:36737ms step_avg:44.64ms
step:824/1555 train_time:36802ms step_avg:44.66ms
step:825/1555 train_time:36859ms step_avg:44.68ms
step:826/1555 train_time:36924ms step_avg:44.70ms
step:827/1555 train_time:36982ms step_avg:44.72ms
step:828/1555 train_time:37046ms step_avg:44.74ms
step:829/1555 train_time:37103ms step_avg:44.76ms
step:830/1555 train_time:37168ms step_avg:44.78ms
step:831/1555 train_time:37224ms step_avg:44.79ms
step:832/1555 train_time:37289ms step_avg:44.82ms
step:833/1555 train_time:37347ms step_avg:44.83ms
step:834/1555 train_time:37411ms step_avg:44.86ms
step:835/1555 train_time:37469ms step_avg:44.87ms
step:836/1555 train_time:37532ms step_avg:44.89ms
step:837/1555 train_time:37591ms step_avg:44.91ms
step:838/1555 train_time:37656ms step_avg:44.93ms
step:839/1555 train_time:37713ms step_avg:44.95ms
step:840/1555 train_time:37779ms step_avg:44.97ms
step:841/1555 train_time:37836ms step_avg:44.99ms
step:842/1555 train_time:37902ms step_avg:45.01ms
step:843/1555 train_time:37959ms step_avg:45.03ms
step:844/1555 train_time:38023ms step_avg:45.05ms
step:845/1555 train_time:38081ms step_avg:45.07ms
step:846/1555 train_time:38145ms step_avg:45.09ms
step:847/1555 train_time:38203ms step_avg:45.10ms
step:848/1555 train_time:38269ms step_avg:45.13ms
step:849/1555 train_time:38327ms step_avg:45.14ms
step:850/1555 train_time:38390ms step_avg:45.16ms
step:851/1555 train_time:38448ms step_avg:45.18ms
step:852/1555 train_time:38512ms step_avg:45.20ms
step:853/1555 train_time:38570ms step_avg:45.22ms
step:854/1555 train_time:38636ms step_avg:45.24ms
step:855/1555 train_time:38692ms step_avg:45.25ms
step:856/1555 train_time:38756ms step_avg:45.28ms
step:857/1555 train_time:38814ms step_avg:45.29ms
step:858/1555 train_time:38879ms step_avg:45.31ms
step:859/1555 train_time:38938ms step_avg:45.33ms
step:860/1555 train_time:39002ms step_avg:45.35ms
step:861/1555 train_time:39059ms step_avg:45.36ms
step:862/1555 train_time:39123ms step_avg:45.39ms
step:863/1555 train_time:39181ms step_avg:45.40ms
step:864/1555 train_time:39246ms step_avg:45.42ms
step:865/1555 train_time:39305ms step_avg:45.44ms
step:866/1555 train_time:39369ms step_avg:45.46ms
step:867/1555 train_time:39427ms step_avg:45.47ms
step:868/1555 train_time:39491ms step_avg:45.50ms
step:869/1555 train_time:39549ms step_avg:45.51ms
step:870/1555 train_time:39614ms step_avg:45.53ms
step:871/1555 train_time:39671ms step_avg:45.55ms
step:872/1555 train_time:39736ms step_avg:45.57ms
step:873/1555 train_time:39792ms step_avg:45.58ms
step:874/1555 train_time:39857ms step_avg:45.60ms
step:875/1555 train_time:39914ms step_avg:45.62ms
step:876/1555 train_time:39979ms step_avg:45.64ms
step:877/1555 train_time:40037ms step_avg:45.65ms
step:878/1555 train_time:40101ms step_avg:45.67ms
step:879/1555 train_time:40159ms step_avg:45.69ms
step:880/1555 train_time:40224ms step_avg:45.71ms
step:881/1555 train_time:40282ms step_avg:45.72ms
step:882/1555 train_time:40346ms step_avg:45.74ms
step:883/1555 train_time:40404ms step_avg:45.76ms
step:884/1555 train_time:40470ms step_avg:45.78ms
step:885/1555 train_time:40528ms step_avg:45.79ms
step:886/1555 train_time:40592ms step_avg:45.82ms
step:887/1555 train_time:40650ms step_avg:45.83ms
step:888/1555 train_time:40714ms step_avg:45.85ms
step:889/1555 train_time:40771ms step_avg:45.86ms
step:890/1555 train_time:40835ms step_avg:45.88ms
step:891/1555 train_time:40893ms step_avg:45.90ms
step:892/1555 train_time:40957ms step_avg:45.92ms
step:893/1555 train_time:41014ms step_avg:45.93ms
step:894/1555 train_time:41079ms step_avg:45.95ms
step:895/1555 train_time:41137ms step_avg:45.96ms
step:896/1555 train_time:41202ms step_avg:45.98ms
step:897/1555 train_time:41260ms step_avg:46.00ms
step:898/1555 train_time:41324ms step_avg:46.02ms
step:899/1555 train_time:41381ms step_avg:46.03ms
step:900/1555 train_time:41446ms step_avg:46.05ms
step:901/1555 train_time:41504ms step_avg:46.06ms
step:902/1555 train_time:41569ms step_avg:46.09ms
step:903/1555 train_time:41627ms step_avg:46.10ms
step:904/1555 train_time:41691ms step_avg:46.12ms
step:905/1555 train_time:41750ms step_avg:46.13ms
step:906/1555 train_time:41813ms step_avg:46.15ms
step:907/1555 train_time:41872ms step_avg:46.17ms
step:908/1555 train_time:41934ms step_avg:46.18ms
step:909/1555 train_time:41992ms step_avg:46.20ms
step:910/1555 train_time:42056ms step_avg:46.22ms
step:911/1555 train_time:42113ms step_avg:46.23ms
step:912/1555 train_time:42178ms step_avg:46.25ms
step:913/1555 train_time:42236ms step_avg:46.26ms
step:914/1555 train_time:42301ms step_avg:46.28ms
step:915/1555 train_time:42359ms step_avg:46.29ms
step:916/1555 train_time:42424ms step_avg:46.31ms
step:917/1555 train_time:42481ms step_avg:46.33ms
step:918/1555 train_time:42546ms step_avg:46.35ms
step:919/1555 train_time:42604ms step_avg:46.36ms
step:920/1555 train_time:42668ms step_avg:46.38ms
step:921/1555 train_time:42726ms step_avg:46.39ms
step:922/1555 train_time:42791ms step_avg:46.41ms
step:923/1555 train_time:42850ms step_avg:46.42ms
step:924/1555 train_time:42913ms step_avg:46.44ms
step:925/1555 train_time:42971ms step_avg:46.46ms
step:926/1555 train_time:43034ms step_avg:46.47ms
step:927/1555 train_time:43092ms step_avg:46.49ms
step:928/1555 train_time:43156ms step_avg:46.50ms
step:929/1555 train_time:43213ms step_avg:46.52ms
step:930/1555 train_time:43279ms step_avg:46.54ms
step:931/1555 train_time:43337ms step_avg:46.55ms
step:932/1555 train_time:43401ms step_avg:46.57ms
step:933/1555 train_time:43459ms step_avg:46.58ms
step:934/1555 train_time:43524ms step_avg:46.60ms
step:935/1555 train_time:43582ms step_avg:46.61ms
step:936/1555 train_time:43647ms step_avg:46.63ms
step:937/1555 train_time:43705ms step_avg:46.64ms
step:938/1555 train_time:43770ms step_avg:46.66ms
step:939/1555 train_time:43828ms step_avg:46.68ms
step:940/1555 train_time:43892ms step_avg:46.69ms
step:941/1555 train_time:43950ms step_avg:46.71ms
step:942/1555 train_time:44014ms step_avg:46.72ms
step:943/1555 train_time:44071ms step_avg:46.74ms
step:944/1555 train_time:44135ms step_avg:46.75ms
step:945/1555 train_time:44192ms step_avg:46.76ms
step:946/1555 train_time:44257ms step_avg:46.78ms
step:947/1555 train_time:44314ms step_avg:46.79ms
step:948/1555 train_time:44380ms step_avg:46.81ms
step:949/1555 train_time:44437ms step_avg:46.83ms
step:950/1555 train_time:44502ms step_avg:46.84ms
step:951/1555 train_time:44560ms step_avg:46.86ms
step:952/1555 train_time:44625ms step_avg:46.87ms
step:953/1555 train_time:44682ms step_avg:46.89ms
step:954/1555 train_time:44747ms step_avg:46.90ms
step:955/1555 train_time:44804ms step_avg:46.92ms
step:956/1555 train_time:44869ms step_avg:46.93ms
step:957/1555 train_time:44927ms step_avg:46.95ms
step:958/1555 train_time:44991ms step_avg:46.96ms
step:959/1555 train_time:45049ms step_avg:46.98ms
step:960/1555 train_time:45113ms step_avg:46.99ms
step:961/1555 train_time:45171ms step_avg:47.00ms
step:962/1555 train_time:45235ms step_avg:47.02ms
step:963/1555 train_time:45292ms step_avg:47.03ms
step:964/1555 train_time:45357ms step_avg:47.05ms
step:965/1555 train_time:45414ms step_avg:47.06ms
step:966/1555 train_time:45480ms step_avg:47.08ms
step:967/1555 train_time:45538ms step_avg:47.09ms
step:968/1555 train_time:45602ms step_avg:47.11ms
step:969/1555 train_time:45660ms step_avg:47.12ms
step:970/1555 train_time:45724ms step_avg:47.14ms
step:971/1555 train_time:45781ms step_avg:47.15ms
step:972/1555 train_time:45847ms step_avg:47.17ms
step:973/1555 train_time:45905ms step_avg:47.18ms
step:974/1555 train_time:45969ms step_avg:47.20ms
step:975/1555 train_time:46027ms step_avg:47.21ms
step:976/1555 train_time:46091ms step_avg:47.22ms
step:977/1555 train_time:46149ms step_avg:47.24ms
step:978/1555 train_time:46212ms step_avg:47.25ms
step:979/1555 train_time:46269ms step_avg:47.26ms
step:980/1555 train_time:46334ms step_avg:47.28ms
step:981/1555 train_time:46392ms step_avg:47.29ms
step:982/1555 train_time:46457ms step_avg:47.31ms
step:983/1555 train_time:46514ms step_avg:47.32ms
step:984/1555 train_time:46579ms step_avg:47.34ms
step:985/1555 train_time:46638ms step_avg:47.35ms
step:986/1555 train_time:46702ms step_avg:47.37ms
step:987/1555 train_time:46760ms step_avg:47.38ms
step:988/1555 train_time:46824ms step_avg:47.39ms
step:989/1555 train_time:46883ms step_avg:47.40ms
step:990/1555 train_time:46947ms step_avg:47.42ms
step:991/1555 train_time:47006ms step_avg:47.43ms
step:992/1555 train_time:47070ms step_avg:47.45ms
step:993/1555 train_time:47127ms step_avg:47.46ms
step:994/1555 train_time:47191ms step_avg:47.48ms
step:995/1555 train_time:47250ms step_avg:47.49ms
step:996/1555 train_time:47315ms step_avg:47.51ms
step:997/1555 train_time:47371ms step_avg:47.51ms
step:998/1555 train_time:47435ms step_avg:47.53ms
step:999/1555 train_time:47493ms step_avg:47.54ms
step:1000/1555 train_time:47557ms step_avg:47.56ms
step:1000/1555 val_loss:3.5774 train_time:47640ms step_avg:47.64ms
step:1001/1555 train_time:47658ms step_avg:47.61ms
step:1002/1555 train_time:47681ms step_avg:47.59ms
step:1003/1555 train_time:47740ms step_avg:47.60ms
step:1004/1555 train_time:47807ms step_avg:47.62ms
step:1005/1555 train_time:47866ms step_avg:47.63ms
step:1006/1555 train_time:47930ms step_avg:47.64ms
step:1007/1555 train_time:47987ms step_avg:47.65ms
step:1008/1555 train_time:48051ms step_avg:47.67ms
step:1009/1555 train_time:48107ms step_avg:47.68ms
step:1010/1555 train_time:48171ms step_avg:47.69ms
step:1011/1555 train_time:48231ms step_avg:47.71ms
step:1012/1555 train_time:48316ms step_avg:47.74ms
step:1013/1555 train_time:48398ms step_avg:47.78ms
step:1014/1555 train_time:48488ms step_avg:47.82ms
step:1015/1555 train_time:48571ms step_avg:47.85ms
step:1016/1555 train_time:48663ms step_avg:47.90ms
step:1017/1555 train_time:48752ms step_avg:47.94ms
step:1018/1555 train_time:48843ms step_avg:47.98ms
step:1019/1555 train_time:48928ms step_avg:48.02ms
step:1020/1555 train_time:49018ms step_avg:48.06ms
step:1021/1555 train_time:49102ms step_avg:48.09ms
step:1022/1555 train_time:49191ms step_avg:48.13ms
step:1023/1555 train_time:49274ms step_avg:48.17ms
step:1024/1555 train_time:49363ms step_avg:48.21ms
step:1025/1555 train_time:49446ms step_avg:48.24ms
step:1026/1555 train_time:49536ms step_avg:48.28ms
step:1027/1555 train_time:49620ms step_avg:48.32ms
step:1028/1555 train_time:49712ms step_avg:48.36ms
step:1029/1555 train_time:49795ms step_avg:48.39ms
step:1030/1555 train_time:49887ms step_avg:48.43ms
step:1031/1555 train_time:49973ms step_avg:48.47ms
step:1032/1555 train_time:50063ms step_avg:48.51ms
step:1033/1555 train_time:50147ms step_avg:48.54ms
step:1034/1555 train_time:50236ms step_avg:48.58ms
step:1035/1555 train_time:50319ms step_avg:48.62ms
step:1036/1555 train_time:50409ms step_avg:48.66ms
step:1037/1555 train_time:50493ms step_avg:48.69ms
step:1038/1555 train_time:50583ms step_avg:48.73ms
step:1039/1555 train_time:50668ms step_avg:48.77ms
step:1040/1555 train_time:50759ms step_avg:48.81ms
step:1041/1555 train_time:50845ms step_avg:48.84ms
step:1042/1555 train_time:50937ms step_avg:48.88ms
step:1043/1555 train_time:51020ms step_avg:48.92ms
step:1044/1555 train_time:51111ms step_avg:48.96ms
step:1045/1555 train_time:51194ms step_avg:48.99ms
step:1046/1555 train_time:51284ms step_avg:49.03ms
step:1047/1555 train_time:51368ms step_avg:49.06ms
step:1048/1555 train_time:51457ms step_avg:49.10ms
step:1049/1555 train_time:51541ms step_avg:49.13ms
step:1050/1555 train_time:51633ms step_avg:49.17ms
step:1051/1555 train_time:51717ms step_avg:49.21ms
step:1052/1555 train_time:51810ms step_avg:49.25ms
step:1053/1555 train_time:51893ms step_avg:49.28ms
step:1054/1555 train_time:51983ms step_avg:49.32ms
step:1055/1555 train_time:52066ms step_avg:49.35ms
step:1056/1555 train_time:52156ms step_avg:49.39ms
step:1057/1555 train_time:52241ms step_avg:49.42ms
step:1058/1555 train_time:52332ms step_avg:49.46ms
step:1059/1555 train_time:52414ms step_avg:49.49ms
step:1060/1555 train_time:52504ms step_avg:49.53ms
step:1061/1555 train_time:52588ms step_avg:49.56ms
step:1062/1555 train_time:52677ms step_avg:49.60ms
step:1063/1555 train_time:52762ms step_avg:49.64ms
step:1064/1555 train_time:52853ms step_avg:49.67ms
step:1065/1555 train_time:52936ms step_avg:49.71ms
step:1066/1555 train_time:53027ms step_avg:49.74ms
step:1067/1555 train_time:53112ms step_avg:49.78ms
step:1068/1555 train_time:53201ms step_avg:49.81ms
step:1069/1555 train_time:53284ms step_avg:49.85ms
step:1070/1555 train_time:53374ms step_avg:49.88ms
step:1071/1555 train_time:53458ms step_avg:49.91ms
step:1072/1555 train_time:53548ms step_avg:49.95ms
step:1073/1555 train_time:53633ms step_avg:49.98ms
step:1074/1555 train_time:53723ms step_avg:50.02ms
step:1075/1555 train_time:53807ms step_avg:50.05ms
step:1076/1555 train_time:53896ms step_avg:50.09ms
step:1077/1555 train_time:53980ms step_avg:50.12ms
step:1078/1555 train_time:54072ms step_avg:50.16ms
step:1079/1555 train_time:54156ms step_avg:50.19ms
step:1080/1555 train_time:54246ms step_avg:50.23ms
step:1081/1555 train_time:54330ms step_avg:50.26ms
step:1082/1555 train_time:54419ms step_avg:50.29ms
step:1083/1555 train_time:54503ms step_avg:50.33ms
step:1084/1555 train_time:54592ms step_avg:50.36ms
step:1085/1555 train_time:54676ms step_avg:50.39ms
step:1086/1555 train_time:54767ms step_avg:50.43ms
step:1087/1555 train_time:54851ms step_avg:50.46ms
step:1088/1555 train_time:54940ms step_avg:50.50ms
step:1089/1555 train_time:55025ms step_avg:50.53ms
step:1090/1555 train_time:55114ms step_avg:50.56ms
step:1091/1555 train_time:55198ms step_avg:50.59ms
step:1092/1555 train_time:55287ms step_avg:50.63ms
step:1093/1555 train_time:55371ms step_avg:50.66ms
step:1094/1555 train_time:55461ms step_avg:50.70ms
step:1095/1555 train_time:55547ms step_avg:50.73ms
step:1096/1555 train_time:55636ms step_avg:50.76ms
step:1097/1555 train_time:55720ms step_avg:50.79ms
step:1098/1555 train_time:55812ms step_avg:50.83ms
step:1099/1555 train_time:55894ms step_avg:50.86ms
step:1100/1555 train_time:55985ms step_avg:50.90ms
step:1101/1555 train_time:56070ms step_avg:50.93ms
step:1102/1555 train_time:56160ms step_avg:50.96ms
step:1103/1555 train_time:56244ms step_avg:50.99ms
step:1104/1555 train_time:56334ms step_avg:51.03ms
step:1105/1555 train_time:56418ms step_avg:51.06ms
step:1106/1555 train_time:56510ms step_avg:51.09ms
step:1107/1555 train_time:56593ms step_avg:51.12ms
step:1108/1555 train_time:56684ms step_avg:51.16ms
step:1109/1555 train_time:56768ms step_avg:51.19ms
step:1110/1555 train_time:56857ms step_avg:51.22ms
step:1111/1555 train_time:56941ms step_avg:51.25ms
step:1112/1555 train_time:57031ms step_avg:51.29ms
step:1113/1555 train_time:57114ms step_avg:51.32ms
step:1114/1555 train_time:57204ms step_avg:51.35ms
step:1115/1555 train_time:57287ms step_avg:51.38ms
step:1116/1555 train_time:57377ms step_avg:51.41ms
step:1117/1555 train_time:57461ms step_avg:51.44ms
step:1118/1555 train_time:57552ms step_avg:51.48ms
step:1119/1555 train_time:57636ms step_avg:51.51ms
step:1120/1555 train_time:57725ms step_avg:51.54ms
step:1121/1555 train_time:57810ms step_avg:51.57ms
step:1122/1555 train_time:57899ms step_avg:51.60ms
step:1123/1555 train_time:57984ms step_avg:51.63ms
step:1124/1555 train_time:58073ms step_avg:51.67ms
step:1125/1555 train_time:58157ms step_avg:51.69ms
step:1126/1555 train_time:58248ms step_avg:51.73ms
step:1127/1555 train_time:58331ms step_avg:51.76ms
step:1128/1555 train_time:58422ms step_avg:51.79ms
step:1129/1555 train_time:58507ms step_avg:51.82ms
step:1130/1555 train_time:58596ms step_avg:51.85ms
step:1131/1555 train_time:58679ms step_avg:51.88ms
step:1132/1555 train_time:58770ms step_avg:51.92ms
step:1133/1555 train_time:58854ms step_avg:51.94ms
step:1134/1555 train_time:58944ms step_avg:51.98ms
step:1135/1555 train_time:59029ms step_avg:52.01ms
step:1136/1555 train_time:59118ms step_avg:52.04ms
step:1137/1555 train_time:59201ms step_avg:52.07ms
step:1138/1555 train_time:59291ms step_avg:52.10ms
step:1139/1555 train_time:59376ms step_avg:52.13ms
step:1140/1555 train_time:59466ms step_avg:52.16ms
step:1141/1555 train_time:59551ms step_avg:52.19ms
step:1142/1555 train_time:59640ms step_avg:52.22ms
step:1143/1555 train_time:59725ms step_avg:52.25ms
step:1144/1555 train_time:59815ms step_avg:52.29ms
step:1145/1555 train_time:59899ms step_avg:52.31ms
step:1146/1555 train_time:59989ms step_avg:52.35ms
step:1147/1555 train_time:60073ms step_avg:52.37ms
step:1148/1555 train_time:60163ms step_avg:52.41ms
step:1149/1555 train_time:60248ms step_avg:52.44ms
step:1150/1555 train_time:60337ms step_avg:52.47ms
step:1151/1555 train_time:60421ms step_avg:52.49ms
step:1152/1555 train_time:60513ms step_avg:52.53ms
step:1153/1555 train_time:60596ms step_avg:52.55ms
step:1154/1555 train_time:60686ms step_avg:52.59ms
step:1155/1555 train_time:60770ms step_avg:52.61ms
step:1156/1555 train_time:60859ms step_avg:52.65ms
step:1157/1555 train_time:60944ms step_avg:52.67ms
step:1158/1555 train_time:61035ms step_avg:52.71ms
step:1159/1555 train_time:61120ms step_avg:52.73ms
step:1160/1555 train_time:61211ms step_avg:52.77ms
step:1161/1555 train_time:61294ms step_avg:52.79ms
step:1162/1555 train_time:61383ms step_avg:52.83ms
step:1163/1555 train_time:61468ms step_avg:52.85ms
step:1164/1555 train_time:61557ms step_avg:52.88ms
step:1165/1555 train_time:61641ms step_avg:52.91ms
step:1166/1555 train_time:61733ms step_avg:52.94ms
step:1167/1555 train_time:61816ms step_avg:52.97ms
step:1168/1555 train_time:61907ms step_avg:53.00ms
step:1169/1555 train_time:61990ms step_avg:53.03ms
step:1170/1555 train_time:62080ms step_avg:53.06ms
step:1171/1555 train_time:62165ms step_avg:53.09ms
step:1172/1555 train_time:62255ms step_avg:53.12ms
step:1173/1555 train_time:62338ms step_avg:53.14ms
step:1174/1555 train_time:62429ms step_avg:53.18ms
step:1175/1555 train_time:62513ms step_avg:53.20ms
step:1176/1555 train_time:62603ms step_avg:53.23ms
step:1177/1555 train_time:62687ms step_avg:53.26ms
step:1178/1555 train_time:62777ms step_avg:53.29ms
step:1179/1555 train_time:62862ms step_avg:53.32ms
step:1180/1555 train_time:62953ms step_avg:53.35ms
step:1181/1555 train_time:63037ms step_avg:53.38ms
step:1182/1555 train_time:63127ms step_avg:53.41ms
step:1183/1555 train_time:63211ms step_avg:53.43ms
step:1184/1555 train_time:63300ms step_avg:53.46ms
step:1185/1555 train_time:63385ms step_avg:53.49ms
step:1186/1555 train_time:63474ms step_avg:53.52ms
step:1187/1555 train_time:63558ms step_avg:53.54ms
step:1188/1555 train_time:63648ms step_avg:53.58ms
step:1189/1555 train_time:63733ms step_avg:53.60ms
step:1190/1555 train_time:63823ms step_avg:53.63ms
step:1191/1555 train_time:63908ms step_avg:53.66ms
step:1192/1555 train_time:63997ms step_avg:53.69ms
step:1193/1555 train_time:64083ms step_avg:53.72ms
step:1194/1555 train_time:64173ms step_avg:53.75ms
step:1195/1555 train_time:64256ms step_avg:53.77ms
step:1196/1555 train_time:64346ms step_avg:53.80ms
step:1197/1555 train_time:64431ms step_avg:53.83ms
step:1198/1555 train_time:64521ms step_avg:53.86ms
step:1199/1555 train_time:64606ms step_avg:53.88ms
step:1200/1555 train_time:64696ms step_avg:53.91ms
step:1201/1555 train_time:64779ms step_avg:53.94ms
step:1202/1555 train_time:64870ms step_avg:53.97ms
step:1203/1555 train_time:64953ms step_avg:53.99ms
step:1204/1555 train_time:65043ms step_avg:54.02ms
step:1205/1555 train_time:65127ms step_avg:54.05ms
step:1206/1555 train_time:65216ms step_avg:54.08ms
step:1207/1555 train_time:65300ms step_avg:54.10ms
step:1208/1555 train_time:65391ms step_avg:54.13ms
step:1209/1555 train_time:65475ms step_avg:54.16ms
step:1210/1555 train_time:65564ms step_avg:54.19ms
step:1211/1555 train_time:65648ms step_avg:54.21ms
step:1212/1555 train_time:65738ms step_avg:54.24ms
step:1213/1555 train_time:65823ms step_avg:54.26ms
step:1214/1555 train_time:65914ms step_avg:54.30ms
step:1215/1555 train_time:65998ms step_avg:54.32ms
step:1216/1555 train_time:66088ms step_avg:54.35ms
step:1217/1555 train_time:66172ms step_avg:54.37ms
step:1218/1555 train_time:66261ms step_avg:54.40ms
step:1219/1555 train_time:66346ms step_avg:54.43ms
step:1220/1555 train_time:66435ms step_avg:54.46ms
step:1221/1555 train_time:66519ms step_avg:54.48ms
step:1222/1555 train_time:66609ms step_avg:54.51ms
step:1223/1555 train_time:66692ms step_avg:54.53ms
step:1224/1555 train_time:66784ms step_avg:54.56ms
step:1225/1555 train_time:66869ms step_avg:54.59ms
step:1226/1555 train_time:66959ms step_avg:54.62ms
step:1227/1555 train_time:67044ms step_avg:54.64ms
step:1228/1555 train_time:67134ms step_avg:54.67ms
step:1229/1555 train_time:67217ms step_avg:54.69ms
step:1230/1555 train_time:67309ms step_avg:54.72ms
step:1231/1555 train_time:67393ms step_avg:54.75ms
step:1232/1555 train_time:67483ms step_avg:54.77ms
step:1233/1555 train_time:67567ms step_avg:54.80ms
step:1234/1555 train_time:67656ms step_avg:54.83ms
step:1235/1555 train_time:67741ms step_avg:54.85ms
step:1236/1555 train_time:67831ms step_avg:54.88ms
step:1237/1555 train_time:67916ms step_avg:54.90ms
step:1238/1555 train_time:68006ms step_avg:54.93ms
step:1239/1555 train_time:68090ms step_avg:54.96ms
step:1240/1555 train_time:68178ms step_avg:54.98ms
step:1241/1555 train_time:68263ms step_avg:55.01ms
step:1242/1555 train_time:68353ms step_avg:55.03ms
step:1243/1555 train_time:68437ms step_avg:55.06ms
step:1244/1555 train_time:68527ms step_avg:55.09ms
step:1245/1555 train_time:68613ms step_avg:55.11ms
step:1246/1555 train_time:68702ms step_avg:55.14ms
step:1247/1555 train_time:68785ms step_avg:55.16ms
step:1248/1555 train_time:68876ms step_avg:55.19ms
step:1249/1555 train_time:68959ms step_avg:55.21ms
step:1250/1555 train_time:69050ms step_avg:55.24ms
step:1250/1555 val_loss:3.4009 train_time:69163ms step_avg:55.33ms
step:1251/1555 train_time:69183ms step_avg:55.30ms
step:1252/1555 train_time:69224ms step_avg:55.29ms
step:1253/1555 train_time:69310ms step_avg:55.32ms
step:1254/1555 train_time:69404ms step_avg:55.35ms
step:1255/1555 train_time:69489ms step_avg:55.37ms
step:1256/1555 train_time:69578ms step_avg:55.40ms
step:1257/1555 train_time:69662ms step_avg:55.42ms
step:1258/1555 train_time:69750ms step_avg:55.45ms
step:1259/1555 train_time:69834ms step_avg:55.47ms
step:1260/1555 train_time:69923ms step_avg:55.49ms
step:1261/1555 train_time:70005ms step_avg:55.52ms
step:1262/1555 train_time:70095ms step_avg:55.54ms
step:1263/1555 train_time:70182ms step_avg:55.57ms
step:1264/1555 train_time:70273ms step_avg:55.60ms
step:1265/1555 train_time:70361ms step_avg:55.62ms
step:1266/1555 train_time:70451ms step_avg:55.65ms
step:1267/1555 train_time:70536ms step_avg:55.67ms
step:1268/1555 train_time:70626ms step_avg:55.70ms
step:1269/1555 train_time:70708ms step_avg:55.72ms
step:1270/1555 train_time:70797ms step_avg:55.75ms
step:1271/1555 train_time:70881ms step_avg:55.77ms
step:1272/1555 train_time:70969ms step_avg:55.79ms
step:1273/1555 train_time:71053ms step_avg:55.82ms
step:1274/1555 train_time:71144ms step_avg:55.84ms
step:1275/1555 train_time:71228ms step_avg:55.87ms
step:1276/1555 train_time:71319ms step_avg:55.89ms
step:1277/1555 train_time:71404ms step_avg:55.92ms
step:1278/1555 train_time:71494ms step_avg:55.94ms
step:1279/1555 train_time:71579ms step_avg:55.96ms
step:1280/1555 train_time:71668ms step_avg:55.99ms
step:1281/1555 train_time:71751ms step_avg:56.01ms
step:1282/1555 train_time:71840ms step_avg:56.04ms
step:1283/1555 train_time:71924ms step_avg:56.06ms
step:1284/1555 train_time:72013ms step_avg:56.08ms
step:1285/1555 train_time:72097ms step_avg:56.11ms
step:1286/1555 train_time:72188ms step_avg:56.13ms
step:1287/1555 train_time:72272ms step_avg:56.16ms
step:1288/1555 train_time:72364ms step_avg:56.18ms
step:1289/1555 train_time:72448ms step_avg:56.20ms
step:1290/1555 train_time:72538ms step_avg:56.23ms
step:1291/1555 train_time:72622ms step_avg:56.25ms
step:1292/1555 train_time:72711ms step_avg:56.28ms
step:1293/1555 train_time:72795ms step_avg:56.30ms
step:1294/1555 train_time:72885ms step_avg:56.33ms
step:1295/1555 train_time:72968ms step_avg:56.35ms
step:1296/1555 train_time:73059ms step_avg:56.37ms
step:1297/1555 train_time:73142ms step_avg:56.39ms
step:1298/1555 train_time:73233ms step_avg:56.42ms
step:1299/1555 train_time:73319ms step_avg:56.44ms
step:1300/1555 train_time:73408ms step_avg:56.47ms
step:1301/1555 train_time:73493ms step_avg:56.49ms
step:1302/1555 train_time:73583ms step_avg:56.52ms
step:1303/1555 train_time:73667ms step_avg:56.54ms
step:1304/1555 train_time:73756ms step_avg:56.56ms
step:1305/1555 train_time:73841ms step_avg:56.58ms
step:1306/1555 train_time:73929ms step_avg:56.61ms
step:1307/1555 train_time:74014ms step_avg:56.63ms
step:1308/1555 train_time:74104ms step_avg:56.65ms
step:1309/1555 train_time:74188ms step_avg:56.68ms
step:1310/1555 train_time:74278ms step_avg:56.70ms
step:1311/1555 train_time:74364ms step_avg:56.72ms
step:1312/1555 train_time:74454ms step_avg:56.75ms
step:1313/1555 train_time:74539ms step_avg:56.77ms
step:1314/1555 train_time:74629ms step_avg:56.79ms
step:1315/1555 train_time:74713ms step_avg:56.82ms
step:1316/1555 train_time:74803ms step_avg:56.84ms
step:1317/1555 train_time:74887ms step_avg:56.86ms
step:1318/1555 train_time:74976ms step_avg:56.89ms
step:1319/1555 train_time:75060ms step_avg:56.91ms
step:1320/1555 train_time:75149ms step_avg:56.93ms
step:1321/1555 train_time:75234ms step_avg:56.95ms
step:1322/1555 train_time:75325ms step_avg:56.98ms
step:1323/1555 train_time:75408ms step_avg:57.00ms
step:1324/1555 train_time:75499ms step_avg:57.02ms
step:1325/1555 train_time:75583ms step_avg:57.04ms
step:1326/1555 train_time:75674ms step_avg:57.07ms
step:1327/1555 train_time:75757ms step_avg:57.09ms
step:1328/1555 train_time:75848ms step_avg:57.11ms
step:1329/1555 train_time:75931ms step_avg:57.13ms
step:1330/1555 train_time:76022ms step_avg:57.16ms
step:1331/1555 train_time:76105ms step_avg:57.18ms
step:1332/1555 train_time:76194ms step_avg:57.20ms
step:1333/1555 train_time:76278ms step_avg:57.22ms
step:1334/1555 train_time:76368ms step_avg:57.25ms
step:1335/1555 train_time:76453ms step_avg:57.27ms
step:1336/1555 train_time:76544ms step_avg:57.29ms
step:1337/1555 train_time:76627ms step_avg:57.31ms
step:1338/1555 train_time:76717ms step_avg:57.34ms
step:1339/1555 train_time:76801ms step_avg:57.36ms
step:1340/1555 train_time:76890ms step_avg:57.38ms
step:1341/1555 train_time:76975ms step_avg:57.40ms
step:1342/1555 train_time:77065ms step_avg:57.43ms
step:1343/1555 train_time:77148ms step_avg:57.44ms
step:1344/1555 train_time:77238ms step_avg:57.47ms
step:1345/1555 train_time:77323ms step_avg:57.49ms
step:1346/1555 train_time:77412ms step_avg:57.51ms
step:1347/1555 train_time:77499ms step_avg:57.53ms
step:1348/1555 train_time:77588ms step_avg:57.56ms
step:1349/1555 train_time:77671ms step_avg:57.58ms
step:1350/1555 train_time:77762ms step_avg:57.60ms
step:1351/1555 train_time:77846ms step_avg:57.62ms
step:1352/1555 train_time:77936ms step_avg:57.64ms
step:1353/1555 train_time:78020ms step_avg:57.66ms
step:1354/1555 train_time:78109ms step_avg:57.69ms
step:1355/1555 train_time:78193ms step_avg:57.71ms
step:1356/1555 train_time:78284ms step_avg:57.73ms
step:1357/1555 train_time:78368ms step_avg:57.75ms
step:1358/1555 train_time:78459ms step_avg:57.78ms
step:1359/1555 train_time:78543ms step_avg:57.79ms
step:1360/1555 train_time:78633ms step_avg:57.82ms
step:1361/1555 train_time:78717ms step_avg:57.84ms
step:1362/1555 train_time:78806ms step_avg:57.86ms
step:1363/1555 train_time:78890ms step_avg:57.88ms
step:1364/1555 train_time:78981ms step_avg:57.90ms
step:1365/1555 train_time:79064ms step_avg:57.92ms
step:1366/1555 train_time:79154ms step_avg:57.95ms
step:1367/1555 train_time:79239ms step_avg:57.97ms
step:1368/1555 train_time:79328ms step_avg:57.99ms
step:1369/1555 train_time:79411ms step_avg:58.01ms
step:1370/1555 train_time:79501ms step_avg:58.03ms
step:1371/1555 train_time:79585ms step_avg:58.05ms
step:1372/1555 train_time:79676ms step_avg:58.07ms
step:1373/1555 train_time:79760ms step_avg:58.09ms
step:1374/1555 train_time:79849ms step_avg:58.11ms
step:1375/1555 train_time:79934ms step_avg:58.13ms
step:1376/1555 train_time:80025ms step_avg:58.16ms
step:1377/1555 train_time:80107ms step_avg:58.18ms
step:1378/1555 train_time:80198ms step_avg:58.20ms
step:1379/1555 train_time:80282ms step_avg:58.22ms
step:1380/1555 train_time:80372ms step_avg:58.24ms
step:1381/1555 train_time:80456ms step_avg:58.26ms
step:1382/1555 train_time:80546ms step_avg:58.28ms
step:1383/1555 train_time:80630ms step_avg:58.30ms
step:1384/1555 train_time:80721ms step_avg:58.32ms
step:1385/1555 train_time:80805ms step_avg:58.34ms
step:1386/1555 train_time:80893ms step_avg:58.36ms
step:1387/1555 train_time:80978ms step_avg:58.38ms
step:1388/1555 train_time:81068ms step_avg:58.41ms
step:1389/1555 train_time:81152ms step_avg:58.42ms
step:1390/1555 train_time:81244ms step_avg:58.45ms
step:1391/1555 train_time:81327ms step_avg:58.47ms
step:1392/1555 train_time:81417ms step_avg:58.49ms
step:1393/1555 train_time:81501ms step_avg:58.51ms
step:1394/1555 train_time:81591ms step_avg:58.53ms
step:1395/1555 train_time:81675ms step_avg:58.55ms
step:1396/1555 train_time:81765ms step_avg:58.57ms
step:1397/1555 train_time:81849ms step_avg:58.59ms
step:1398/1555 train_time:81939ms step_avg:58.61ms
step:1399/1555 train_time:82024ms step_avg:58.63ms
step:1400/1555 train_time:82113ms step_avg:58.65ms
step:1401/1555 train_time:82198ms step_avg:58.67ms
step:1402/1555 train_time:82287ms step_avg:58.69ms
step:1403/1555 train_time:82370ms step_avg:58.71ms
step:1404/1555 train_time:82462ms step_avg:58.73ms
step:1405/1555 train_time:82546ms step_avg:58.75ms
step:1406/1555 train_time:82636ms step_avg:58.77ms
step:1407/1555 train_time:82720ms step_avg:58.79ms
step:1408/1555 train_time:82810ms step_avg:58.81ms
step:1409/1555 train_time:82894ms step_avg:58.83ms
step:1410/1555 train_time:82985ms step_avg:58.85ms
step:1411/1555 train_time:83068ms step_avg:58.87ms
step:1412/1555 train_time:83160ms step_avg:58.90ms
step:1413/1555 train_time:83244ms step_avg:58.91ms
step:1414/1555 train_time:83333ms step_avg:58.93ms
step:1415/1555 train_time:83417ms step_avg:58.95ms
step:1416/1555 train_time:83506ms step_avg:58.97ms
step:1417/1555 train_time:83591ms step_avg:58.99ms
step:1418/1555 train_time:83680ms step_avg:59.01ms
step:1419/1555 train_time:83765ms step_avg:59.03ms
step:1420/1555 train_time:83856ms step_avg:59.05ms
step:1421/1555 train_time:83940ms step_avg:59.07ms
step:1422/1555 train_time:84029ms step_avg:59.09ms
step:1423/1555 train_time:84112ms step_avg:59.11ms
step:1424/1555 train_time:84203ms step_avg:59.13ms
step:1425/1555 train_time:84287ms step_avg:59.15ms
step:1426/1555 train_time:84377ms step_avg:59.17ms
step:1427/1555 train_time:84461ms step_avg:59.19ms
step:1428/1555 train_time:84551ms step_avg:59.21ms
step:1429/1555 train_time:84636ms step_avg:59.23ms
step:1430/1555 train_time:84726ms step_avg:59.25ms
step:1431/1555 train_time:84809ms step_avg:59.27ms
step:1432/1555 train_time:84900ms step_avg:59.29ms
step:1433/1555 train_time:84984ms step_avg:59.31ms
step:1434/1555 train_time:85074ms step_avg:59.33ms
step:1435/1555 train_time:85159ms step_avg:59.34ms
step:1436/1555 train_time:85249ms step_avg:59.37ms
step:1437/1555 train_time:85333ms step_avg:59.38ms
step:1438/1555 train_time:85425ms step_avg:59.41ms
step:1439/1555 train_time:85507ms step_avg:59.42ms
step:1440/1555 train_time:85597ms step_avg:59.44ms
step:1441/1555 train_time:85682ms step_avg:59.46ms
step:1442/1555 train_time:85772ms step_avg:59.48ms
step:1443/1555 train_time:85856ms step_avg:59.50ms
step:1444/1555 train_time:85947ms step_avg:59.52ms
step:1445/1555 train_time:86031ms step_avg:59.54ms
step:1446/1555 train_time:86123ms step_avg:59.56ms
step:1447/1555 train_time:86205ms step_avg:59.58ms
step:1448/1555 train_time:86295ms step_avg:59.60ms
step:1449/1555 train_time:86379ms step_avg:59.61ms
step:1450/1555 train_time:86468ms step_avg:59.63ms
step:1451/1555 train_time:86553ms step_avg:59.65ms
step:1452/1555 train_time:86645ms step_avg:59.67ms
step:1453/1555 train_time:86728ms step_avg:59.69ms
step:1454/1555 train_time:86818ms step_avg:59.71ms
step:1455/1555 train_time:86902ms step_avg:59.73ms
step:1456/1555 train_time:86991ms step_avg:59.75ms
step:1457/1555 train_time:87075ms step_avg:59.76ms
step:1458/1555 train_time:87166ms step_avg:59.78ms
step:1459/1555 train_time:87249ms step_avg:59.80ms
step:1460/1555 train_time:87339ms step_avg:59.82ms
step:1461/1555 train_time:87423ms step_avg:59.84ms
step:1462/1555 train_time:87511ms step_avg:59.86ms
step:1463/1555 train_time:87596ms step_avg:59.87ms
step:1464/1555 train_time:87686ms step_avg:59.89ms
step:1465/1555 train_time:87770ms step_avg:59.91ms
step:1466/1555 train_time:87861ms step_avg:59.93ms
step:1467/1555 train_time:87945ms step_avg:59.95ms
step:1468/1555 train_time:88035ms step_avg:59.97ms
step:1469/1555 train_time:88119ms step_avg:59.99ms
step:1470/1555 train_time:88208ms step_avg:60.01ms
step:1471/1555 train_time:88292ms step_avg:60.02ms
step:1472/1555 train_time:88382ms step_avg:60.04ms
step:1473/1555 train_time:88467ms step_avg:60.06ms
step:1474/1555 train_time:88556ms step_avg:60.08ms
step:1475/1555 train_time:88642ms step_avg:60.10ms
step:1476/1555 train_time:88731ms step_avg:60.12ms
step:1477/1555 train_time:88815ms step_avg:60.13ms
step:1478/1555 train_time:88905ms step_avg:60.15ms
step:1479/1555 train_time:88989ms step_avg:60.17ms
step:1480/1555 train_time:89079ms step_avg:60.19ms
step:1481/1555 train_time:89164ms step_avg:60.21ms
step:1482/1555 train_time:89253ms step_avg:60.22ms
step:1483/1555 train_time:89338ms step_avg:60.24ms
step:1484/1555 train_time:89428ms step_avg:60.26ms
step:1485/1555 train_time:89512ms step_avg:60.28ms
step:1486/1555 train_time:89601ms step_avg:60.30ms
step:1487/1555 train_time:89685ms step_avg:60.31ms
step:1488/1555 train_time:89776ms step_avg:60.33ms
step:1489/1555 train_time:89861ms step_avg:60.35ms
step:1490/1555 train_time:89952ms step_avg:60.37ms
step:1491/1555 train_time:90036ms step_avg:60.39ms
step:1492/1555 train_time:90128ms step_avg:60.41ms
step:1493/1555 train_time:90211ms step_avg:60.42ms
step:1494/1555 train_time:90301ms step_avg:60.44ms
step:1495/1555 train_time:90385ms step_avg:60.46ms
step:1496/1555 train_time:90474ms step_avg:60.48ms
step:1497/1555 train_time:90560ms step_avg:60.49ms
step:1498/1555 train_time:90649ms step_avg:60.51ms
step:1499/1555 train_time:90733ms step_avg:60.53ms
step:1500/1555 train_time:90824ms step_avg:60.55ms
step:1500/1555 val_loss:3.2972 train_time:90937ms step_avg:60.62ms
step:1501/1555 train_time:90956ms step_avg:60.60ms
step:1502/1555 train_time:90997ms step_avg:60.58ms
step:1503/1555 train_time:91086ms step_avg:60.60ms
step:1504/1555 train_time:91178ms step_avg:60.62ms
step:1505/1555 train_time:91261ms step_avg:60.64ms
step:1506/1555 train_time:91351ms step_avg:60.66ms
step:1507/1555 train_time:91434ms step_avg:60.67ms
step:1508/1555 train_time:91523ms step_avg:60.69ms
step:1509/1555 train_time:91605ms step_avg:60.71ms
step:1510/1555 train_time:91696ms step_avg:60.73ms
step:1511/1555 train_time:91778ms step_avg:60.74ms
step:1512/1555 train_time:91867ms step_avg:60.76ms
step:1513/1555 train_time:91953ms step_avg:60.78ms
step:1514/1555 train_time:92045ms step_avg:60.80ms
step:1515/1555 train_time:92135ms step_avg:60.82ms
step:1516/1555 train_time:92228ms step_avg:60.84ms
step:1517/1555 train_time:92315ms step_avg:60.85ms
step:1518/1555 train_time:92403ms step_avg:60.87ms
step:1519/1555 train_time:92487ms step_avg:60.89ms
step:1520/1555 train_time:92577ms step_avg:60.91ms
step:1521/1555 train_time:92660ms step_avg:60.92ms
step:1522/1555 train_time:92750ms step_avg:60.94ms
step:1523/1555 train_time:92834ms step_avg:60.95ms
step:1524/1555 train_time:92924ms step_avg:60.97ms
step:1525/1555 train_time:93009ms step_avg:60.99ms
step:1526/1555 train_time:93102ms step_avg:61.01ms
step:1527/1555 train_time:93187ms step_avg:61.03ms
step:1528/1555 train_time:93281ms step_avg:61.05ms
step:1529/1555 train_time:93365ms step_avg:61.06ms
step:1530/1555 train_time:93455ms step_avg:61.08ms
step:1531/1555 train_time:93539ms step_avg:61.10ms
step:1532/1555 train_time:93629ms step_avg:61.12ms
step:1533/1555 train_time:93712ms step_avg:61.13ms
step:1534/1555 train_time:93802ms step_avg:61.15ms
step:1535/1555 train_time:93887ms step_avg:61.16ms
step:1536/1555 train_time:93979ms step_avg:61.18ms
step:1537/1555 train_time:94064ms step_avg:61.20ms
step:1538/1555 train_time:94154ms step_avg:61.22ms
step:1539/1555 train_time:94240ms step_avg:61.23ms
step:1540/1555 train_time:94330ms step_avg:61.25ms
step:1541/1555 train_time:94414ms step_avg:61.27ms
step:1542/1555 train_time:94505ms step_avg:61.29ms
step:1543/1555 train_time:94589ms step_avg:61.30ms
step:1544/1555 train_time:94680ms step_avg:61.32ms
step:1545/1555 train_time:94763ms step_avg:61.34ms
step:1546/1555 train_time:94854ms step_avg:61.35ms
step:1547/1555 train_time:94939ms step_avg:61.37ms
step:1548/1555 train_time:95030ms step_avg:61.39ms
step:1549/1555 train_time:95114ms step_avg:61.40ms
step:1550/1555 train_time:95204ms step_avg:61.42ms
step:1551/1555 train_time:95289ms step_avg:61.44ms
step:1552/1555 train_time:95380ms step_avg:61.46ms
step:1553/1555 train_time:95464ms step_avg:61.47ms
step:1554/1555 train_time:95555ms step_avg:61.49ms
step:1555/1555 train_time:95639ms step_avg:61.50ms
step:1555/1555 val_loss:3.2807 train_time:95753ms step_avg:61.58ms
peak memory allocated: 31630 MiB reserved: 46658 MiB