Delete training_original

training_original/binary16_train.py

@@ -1,997 +0,0 @@
-#!/usr/bin/env python3
-
-import argparse
-import contextlib
-import json
-import math
-import os
-import random
-import time
-from pathlib import Path
-
-import numpy as np
-import torch
-import torch.distributed as dist
-from torch.nn.parallel import DistributedDataParallel as DDP
-from transformers import AutoTokenizer
-
-from classic_train import (
-    Logger,
-    TokenShardStore,
-    all_reduce_sum,
-    cleanup_distributed,
-    configure_optimizer,
-    estimate_loss,
-    generate_diagnostics,
-    log_input_representation_diagnostics,
-    get_learning_rate,
-    get_rank,
-    get_world_size,
-    human_tokens,
-    is_main_process,
-    save_model_safetensors,
-    load_checkpoint,
-    save_checkpoint,
-    setup_distributed,
-    timestamp,
-)
-
-from model_16_dim_bin import (
-    Binary16Config,
-    Binary16ForCausalLM,
-)
-
-
-def parse_args():
-    parser = argparse.ArgumentParser()
-
-    # Paths
-    parser.add_argument("--data_dir", required=True)
-    parser.add_argument("--output_dir", required=True)
-    parser.add_argument(
-        "--tokenizer",
-        default="HuggingFaceTB/SmolLM2-135M",
-    )
-    parser.add_argument("--tokenizer_revision", default=None)
-    parser.add_argument("--resume", default=None)
-
-    # Architecture: держать идентично classic run.
-    parser.add_argument("--d_model", type=int, default=960)
-    parser.add_argument("--n_layer", type=int, default=24)
-    parser.add_argument("--n_head", type=int, default=15)
-    parser.add_argument(
-        "--ffn_multiplier",
-        type=float,
-        default=8.0 / 3.0,
-    )
-    parser.add_argument("--multiple_of", type=int, default=256)
-    parser.add_argument("--sequence_length", type=int, default=2048)
-    parser.add_argument("--rope_theta", type=float, default=10000.0)
-    parser.add_argument("--dropout", type=float, default=0.0)
-    parser.add_argument("--rms_norm_eps", type=float, default=1e-5)
-    parser.add_argument(
-        "--gradient_checkpointing",
-        action="store_true",
-    )
-    parser.add_argument("--compile", action="store_true")
-
-    # Binary input
-    parser.add_argument(
-        "--binary_encoding",
-        choices=("zero_one", "bipolar"),
-        default="zero_one",
-        help=(
-            "zero_one reproduces the original 0/1 experiment. "
-            "Do not mix encodings inside the main comparison."
-        ),
-    )
-    parser.add_argument(
-        "--binary_scale",
-        type=float,
-        default=1.0,
-    )
-    parser.add_argument(
-        "--binary_permutation_seed",
-        type=int,
-        default=None,
-        help=(
-            "None gives canonical token-ID bits. A seed creates a fixed "
-            "injective random reassignment of token IDs to 16-bit codes."
-        ),
-    )
-
-    # Training
-    parser.add_argument(
-        "--micro_batch_size",
-        type=int,
-        default=2,
-    )
-    parser.add_argument(
-        "--gradient_accumulation_steps",
-        type=int,
-        default=16,
-    )
-    parser.add_argument("--max_steps", type=int, default=200_000)
-    parser.add_argument(
-        "--max_tokens",
-        type=int,
-        default=0,
-    )
-    parser.add_argument(
-        "--learning_rate",
-        type=float,
-        default=3e-4,
-    )
-    parser.add_argument(
-        "--min_learning_rate",
-        type=float,
-        default=3e-5,
-    )
-    parser.add_argument("--warmup_steps", type=int, default=2000)
-    parser.add_argument("--weight_decay", type=float, default=0.1)
-    parser.add_argument("--beta1", type=float, default=0.9)
-    parser.add_argument("--beta2", type=float, default=0.95)
-    parser.add_argument("--grad_clip", type=float, default=1.0)
-    parser.add_argument("--seed", type=int, default=42)
-
-    # Dataset cache
-    parser.add_argument(
-        "--train_cache_shards",
-        type=int,
-        default=2,
-    )
-    parser.add_argument(
-        "--valid_cache_shards",
-        type=int,
-        default=2,
-    )
-
-    # Diagnostics
-    parser.add_argument(
-        "--diagnostic_interval_seconds",
-        type=float,
-        default=3600.0,
-    )
-    parser.add_argument(
-        "--log_interval_steps",
-        type=int,
-        default=20,
-    )
-    parser.add_argument(
-        "--eval_batches",
-        type=int,
-        default=32,
-    )
-    parser.add_argument(
-        "--eval_batch_size",
-        type=int,
-        default=2,
-    )
-    parser.add_argument(
-        "--generation_tokens",
-        type=int,
-        default=24,
-    )
-    parser.add_argument(
-        "--save_interval_steps",
-        type=int,
-        default=2000,
-    )
-    parser.add_argument(
-        "--batches_per_shard",
-        type=int,
-        default=256,
-        help=(
-            "Number of microbatches sampled from one resident shard "
-            "before loading the next shard."
-        ),
-    )
-    parser.add_argument(
-        "--latest_save_interval_steps",
-        type=int,
-        default=2000,
-    )
-    
-    parser.add_argument(
-        "--milestone_save_interval_steps",
-        type=int,
-        default=20000,
-    )
-    
-    return parser.parse_args()
-
-
-def verify_binary_interface(
-    model: Binary16ForCausalLM,
-    vocab_size: int,
-    device: torch.device,
-):
-    embedding = model.token_embeddings
-
-    if any(True for _ in embedding.parameters()):
-        raise RuntimeError(
-            "FixedBinary16Embedding unexpectedly contains parameters"
-        )
-
-    expected_shape = (vocab_size, 16)
-
-    if tuple(embedding.codebook.shape) != expected_shape:
-        raise RuntimeError(
-            f"Invalid codebook shape: "
-            f"{tuple(embedding.codebook.shape)} != {expected_shape}"
-        )
-
-    if embedding.codebook.requires_grad:
-        raise RuntimeError("Binary codebook must not require gradients")
-
-    # Проверяем первые canonical codes только если permutation выключена.
-    if model.config.binary_permutation_seed is None:
-        expected = torch.tensor(
-            [
-                [0, 0, 0, 0, 0, 0, 0, 0,
-                 0, 0, 0, 0, 0, 0, 0, 0],
-                [1, 0, 0, 0, 0, 0, 0, 0,
-                 0, 0, 0, 0, 0, 0, 0, 0],
-                [0, 1, 0, 0, 0, 0, 0, 0,
-                 0, 0, 0, 0, 0, 0, 0, 0],
-                [1, 1, 0, 0, 0, 0, 0, 0,
-                 0, 0, 0, 0, 0, 0, 0, 0],
-            ],
-            dtype=torch.float32,
-            device=embedding.codebook.device,
-        )
-
-        actual = embedding.codebook[:4]
-
-        if model.config.binary_encoding == "bipolar":
-            expected = expected.mul(2).sub(1)
-
-        if not torch.equal(actual, expected):
-            raise RuntimeError(
-                "Canonical binary code sanity check failed"
-            )
-
-    ids = torch.tensor(
-        [[0, min(1, vocab_size - 1), vocab_size - 1]],
-        device=device,
-        dtype=torch.long,
-    )
-
-    output = embedding(ids)
-
-    expected_output_shape = (
-        1,
-        3,
-        model.config.d_model,
-    )
-
-    if tuple(output.shape) != expected_output_shape:
-        raise RuntimeError(
-            f"Invalid lifted shape: "
-            f"{tuple(output.shape)} != {expected_output_shape}"
-        )
-
-
-def find_nonfinite_gradients(model, max_names=20):
-    bad = []
-
-    for name, parameter in model.named_parameters():
-        gradient = parameter.grad
-
-        if gradient is None:
-            continue
-
-        finite = torch.isfinite(gradient)
-
-        if not bool(finite.all()):
-            bad.append(
-                {
-                    "name": name,
-                    "shape": tuple(gradient.shape),
-                    "dtype": str(gradient.dtype),
-                    "nan": int(torch.isnan(gradient).sum().item()),
-                    "inf": int(torch.isinf(gradient).sum().item()),
-                }
-            )
-
-            if len(bad) >= max_names:
-                break
-
-    return bad
-
-
-def grad_norm_for_named_parameters(named_parameters):
-    squares = []
-
-    for _, parameter in named_parameters:
-        if parameter.grad is None:
-            continue
-
-        gradient = parameter.grad.detach().float()
-        squares.append(gradient.pow(2).sum())
-
-    if not squares:
-        return 0.0
-
-    return torch.sqrt(torch.stack(squares).sum()).item()
-
-
-def main():
-    args = parse_args()
-
-    distributed, local_rank, device = setup_distributed()
-    rank = get_rank()
-    world_size = get_world_size()
-
-    if device.type != "cuda":
-        raise RuntimeError("CUDA is required")
-
-    torch.backends.cuda.matmul.allow_tf32 = True
-    torch.backends.cudnn.allow_tf32 = True
-    torch.backends.cuda.enable_flash_sdp(True)
-    torch.backends.cuda.enable_mem_efficient_sdp(True)
-    torch.backends.cuda.enable_math_sdp(True)
-
-    seed = args.seed + rank
-    random.seed(seed)
-    np.random.seed(seed)
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
-
-    Path(args.output_dir).mkdir(
-        parents=True,
-        exist_ok=True,
-    )
-
-    logger = Logger(
-        os.path.join(args.output_dir, "train.log")
-    )
-
-    tokenizer = AutoTokenizer.from_pretrained(
-        args.tokenizer,
-        revision=args.tokenizer_revision,
-        use_fast=True,
-    )
-
-    vocab_size = len(tokenizer)
-
-    if vocab_size > 65536:
-        raise ValueError(
-            f"Tokenizer length {vocab_size} exceeds 16-bit capacity"
-        )
-
-    if args.d_model % 16 != 0:
-        raise ValueError(
-            "d_model must be divisible by 16 for parameter-free tiling"
-        )
-
-    config = Binary16Config(
-        vocab_size=vocab_size,
-        d_model=args.d_model,
-        n_layer=args.n_layer,
-        n_head=args.n_head,
-        ffn_multiplier=args.ffn_multiplier,
-        multiple_of=args.multiple_of,
-        block_size=args.sequence_length,
-        rope_theta=args.rope_theta,
-        dropout=args.dropout,
-        rms_norm_eps=args.rms_norm_eps,
-        pad_token_id=tokenizer.pad_token_id,
-        bos_token_id=tokenizer.bos_token_id,
-        eos_token_id=tokenizer.eos_token_id,
-        tie_word_embeddings=False,
-        binary_dim=16,
-        binary_encoding=args.binary_encoding,
-        binary_scale=args.binary_scale,
-        binary_permutation_seed=args.binary_permutation_seed,
-    )
-
-    raw_model = Binary16ForCausalLM(config)
-
-    if args.gradient_checkpointing:
-        raw_model.gradient_checkpointing = True
-
-    raw_model.to(device)
-
-    if is_main_process():
-        first_parameter = next(raw_model.parameters())
-    
-        logger.log(
-            "[dtype] "
-            f"parameter_dtype={first_parameter.dtype}, "
-            f"optimizer_master_expected=float32"
-        )
-    
-    if next(raw_model.parameters()).dtype != torch.float32:
-        raise RuntimeError(
-            "Model parameters must remain FP32; "
-            "BF16 should be enabled only through autocast"
-        )
-
-    verify_binary_interface(
-        model=raw_model,
-        vocab_size=vocab_size,
-        device=device,
-    )
-
-    parameter_counts = raw_model.count_parameters()
-
-    if is_main_process():
-        logger.log(
-            "[model] "
-            + ", ".join(
-                f"{key}={value:,}"
-                for key, value in parameter_counts.items()
-            )
-        )
-        logger.log(
-            "[binary_input] "
-            f"encoding={args.binary_encoding}, "
-            f"scale={args.binary_scale}, "
-            f"permutation_seed={args.binary_permutation_seed}, "
-            f"codebook_shape={tuple(raw_model.token_embeddings.codebook.shape)}, "
-            f"codebook_is_parameter=False"
-        )
-        logger.log(
-            "[config] "
-            + json.dumps(
-                config.to_dict(),
-                ensure_ascii=False,
-                sort_keys=True,
-            )
-        )
-        logger.log(
-            f"[run] world_size={world_size}, "
-            f"micro_batch={args.micro_batch_size}, "
-            f"grad_accum={args.gradient_accumulation_steps}, "
-            f"seq={args.sequence_length}, "
-            f"global_tokens_per_step="
-            #f"{world_size * args.micro_batch_size * args.gradient_accumulation_steps * (args.sequence_length - 1):,}"
-            f"{world_size * args.micro_batch_size * args.gradient_accumulation_steps * args.sequence_length:,}"
-        )
-
-    optimizer = configure_optimizer(
-        raw_model,
-        learning_rate=args.learning_rate,
-        weight_decay=args.weight_decay,
-        betas=(args.beta1, args.beta2),
-        fused=True,
-    )
-
-    start_step = 0
-    tokens_seen = 0
-
-    if args.resume is not None:
-        start_step, tokens_seen = load_checkpoint(
-            args.resume,
-            raw_model,
-            optimizer,
-            device,
-        )
-
-        logger.log(
-            f"[resume] path={args.resume}, "
-            f"step={start_step}, "
-            f"tokens_seen={tokens_seen:,}"
-        )
-
-    model = raw_model
-
-    if args.compile:
-        model = torch.compile(
-            model,
-            mode="max-autotune",
-            dynamic=False,
-        )
-
-    if distributed:
-        '''model = DDP(
-            model,
-            device_ids=[local_rank],
-            output_device=local_rank,
-            broadcast_buffers=False,
-            gradient_as_bucket_view=True,
-            static_graph=not args.gradient_checkpointing,
-        )'''
-        model = DDP(
-            model,
-            device_ids=[local_rank],
-            output_device=local_rank,
-            broadcast_buffers=False,
-            gradient_as_bucket_view=True,
-            static_graph=False,
-            find_unused_parameters=False,
-        )
-
-    train_data = TokenShardStore(
-        data_dir=args.data_dir,
-        split="train",
-        sequence_length=args.sequence_length,
-        cache_shards=args.train_cache_shards,
-        seed=args.seed,
-        batches_per_shard=args.batches_per_shard,
-    )
-
-    valid_data = TokenShardStore(
-        data_dir=args.data_dir,
-        split="valid",
-        sequence_length=args.sequence_length,
-        cache_shards=args.valid_cache_shards,
-        seed=args.seed + 10_000,
-        batches_per_shard=max(
-            args.batches_per_shard,
-            args.eval_batches,
-        ),
-    )
-
-    train_generator = torch.Generator(device="cpu")
-    train_generator.manual_seed(
-        args.seed + rank * 100_003
-    )
-
-    prompts = [
-        # English factual completion
-        "London is the capital of",
-        "The capital of France is",
-        "The largest planet in the Solar System is",
-        "Water freezes at",
-        "The chemical symbol for gold is",
-        "The Pacific Ocean is",
-        "The human heart pumps",
-        "The Second World War ended in",
-        "The author of Romeo and Juliet was",
-        "A triangle has",
-    
-        # English continuation and grammar
-        "Once upon a time, there was",
-        "The scientist opened the laboratory door and",
-        "When the rain finally stopped,",
-        "She went to the store because",
-        "If I had known about the problem,",
-        "The old house on the hill",
-        "Although the experiment failed,",
-        "In order to solve this problem, we need to",
-        "The main difference between cats and dogs is",
-        "This article explains how to",
-    
-        # Definitions and explanations
-        "Photosynthesis is the process by which",
-        "Gravity is a force that",
-        "A computer program is",
-        "Democracy can be defined as",
-        "Machine learning is used to",
-        "The purpose of a database is to",
-        "An ecosystem consists of",
-        "Inflation occurs when",
-        "The Internet allows people to",
-        "Energy cannot be created or destroyed, but",
-    
-        # Arithmetic and symbolic patterns
-        "2 + 2 =",
-        "10 - 3 =",
-        "6 * 7 =",
-        "12 / 4 =",
-        "1, 2, 3, 4,",
-        "2, 4, 6, 8,",
-        "The square root of 9 is",
-        "If x = 5, then x + 2 =",
-        "One hundred divided by ten equals",
-        "The next number after 99 is",
-    ]
-
-    interval_loss_sum = 0.0
-    interval_loss_tokens = 0
-    interval_start_tokens = tokens_seen
-    interval_start_time = time.monotonic()
-    last_diagnostic_time = time.monotonic()
-
-    if is_main_process():
-        logger.log(f"[start] {timestamp()}")
-                
-        log_input_representation_diagnostics(
-            raw_model=raw_model,
-            tokenizer=tokenizer,
-            logger=logger,
-            device=device,
-            step=start_step,
-        )
-
-    model.train()
-    optimizer.zero_grad(set_to_none=True)
-
-    for step in range(start_step, args.max_steps):
-        lr = get_learning_rate(
-            step=step,
-            max_steps=args.max_steps,
-            warmup_steps=args.warmup_steps,
-            learning_rate=args.learning_rate,
-            min_learning_rate=args.min_learning_rate,
-        )
-
-        for group in optimizer.param_groups:
-            group["lr"] = lr
-
-        step_loss_sum = 0.0
-        step_loss_tokens = 0
-
-        for micro_step in range(
-            args.gradient_accumulation_steps
-        ):
-            batch = train_data.sample_batch(
-                batch_size=args.micro_batch_size,
-                generator=train_generator,
-            ).to(device, non_blocking=True)
-
-            # forward() сам сдвигает logits и labels на один токен.
-            #input_ids = batch[:, :-1]
-            #labels = input_ids
-            
-            input_ids = batch[:, :-1]
-            labels = batch[:, 1:]
-
-            should_sync = (
-                micro_step
-                == args.gradient_accumulation_steps - 1
-            )
-
-            sync_context = contextlib.nullcontext()
-
-            if distributed and not should_sync:
-                sync_context = model.no_sync()
-
-            with sync_context:
-                with torch.autocast(
-                    device_type="cuda",
-                    dtype=torch.bfloat16,
-                ):
-                    outputs = model(
-                        input_ids=input_ids,
-                        labels=labels,
-                        return_dict=True,
-                    )
-
-                    loss = (
-                        outputs.loss
-                        / args.gradient_accumulation_steps
-                    )
-
-                loss.backward()
-
-            # Первый label не имеет предшествующего logit после внутреннего shift.
-            
-            #local_tokens = labels[:, 1:].numel()
-
-            local_tokens = labels.numel()
-
-            step_loss_sum += (
-                float(outputs.loss.detach())
-                * local_tokens
-            )
-            step_loss_tokens += local_tokens
-            
-        bad_gradients = find_nonfinite_gradients(raw_model)
-        
-        local_bad = torch.tensor(
-            [1 if bad_gradients else 0],
-            device=device,
-            dtype=torch.int32,
-        )
-        
-        if distributed:
-            dist.all_reduce(
-                local_bad,
-                op=dist.ReduceOp.MAX,
-            )
-        
-        if int(local_bad.item()) != 0:
-            if bad_gradients:
-                logger.log(
-                    "[nonfinite_gradients] "
-                    + json.dumps(
-                        bad_gradients,
-                        ensure_ascii=False,
-                    )
-                )
-        
-            emergency_path = os.path.join(
-                args.output_dir,
-                f"checkpoint_nonfinite_step_{step + 1:07d}.pt",
-            )
-        
-            save_checkpoint(
-                path=emergency_path,
-                raw_model=raw_model,
-                optimizer=optimizer,
-                step=step,
-                tokens_seen=tokens_seen,
-                args=args,
-            )
-        
-            raise RuntimeError(
-                f"Non-finite gradients at step {step + 1}"
-            )
-
-        input_norm = grad_norm_for_named_parameters(
-            raw_model.token_embeddings.named_parameters()
-        )
-        
-        body_norm = grad_norm_for_named_parameters(
-            (
-                (name, parameter)
-                for name, parameter in raw_model.named_parameters()
-                if not name.startswith("token_embeddings.")
-                and not name.startswith("lm_head.")
-            )
-        )
-        
-        output_norm = grad_norm_for_named_parameters(
-            raw_model.lm_head.named_parameters()
-        )
-        
-        #logger.log(
-        #    f"[grad_groups] input={input_norm:.3f}, "
-        #    f"body={body_norm:.3f}, output={output_norm:.3f}"
-        #)
-        
-        if args.grad_clip > 0:
-            grad_norm = torch.nn.utils.clip_grad_norm_(
-                raw_model.parameters(),
-                args.grad_clip,
-                error_if_nonfinite=True,
-            )
-        else:
-            grad_norm = torch.tensor(
-                float("nan"),
-                device=device,
-            )
-
-        optimizer.step()
-        optimizer.zero_grad(set_to_none=True)
-
-        global_step_tokens = step_loss_tokens * world_size
-        tokens_seen += global_step_tokens
-
-        loss_stats = torch.tensor(
-            [step_loss_sum, step_loss_tokens],
-            device=device,
-            dtype=torch.float64,
-        )
-        all_reduce_sum(loss_stats)
-
-        global_loss_sum = float(loss_stats[0].item())
-        global_loss_tokens = int(loss_stats[1].item())
-
-        interval_loss_sum += global_loss_sum
-        interval_loss_tokens += global_loss_tokens
-
-        completed_step = step + 1
-
-        if (
-            completed_step % args.log_interval_steps == 0
-            and is_main_process()
-        ):
-            mean_step_loss = (
-                global_loss_sum / global_loss_tokens
-            )
-
-            logger.log(
-                f"step {completed_step:07d}: "
-                f"loss {mean_step_loss:.4f}, "
-                f"lr {lr:.8f}, "
-                f"grad_norm {float(grad_norm):.4f}, "
-                f"tokens_seen={human_tokens(tokens_seen)}, "
-                f"{timestamp()}"
-            )
-
-        now = time.monotonic()
-
-        '''diagnostic_due = (
-            now - last_diagnostic_time
-            >= args.diagnostic_interval_seconds
-        )'''
-
-        diagnostic_due = (
-            completed_step % 700 == 0
-        )
-
-        final_step = completed_step == args.max_steps
-
-        token_budget_reached = (
-            args.max_tokens > 0
-            and tokens_seen >= args.max_tokens
-        )
-
-        if diagnostic_due or final_step or token_budget_reached:
-            if distributed:
-                dist.barrier()
-        
-            elapsed = max(
-                now - interval_start_time,
-                1e-9,
-            )
-
-            interval_tokens = (
-                tokens_seen - interval_start_tokens
-            )
-
-            tokens_per_second = (
-                interval_tokens / elapsed
-            )
-
-            train_loss = (
-                interval_loss_sum
-                / max(interval_loss_tokens, 1)
-            )
-
-            val_loss, val_tokens = estimate_loss(
-                model=model,
-                dataset=valid_data,
-                device=device,
-                batch_size=args.eval_batch_size,
-                eval_batches=args.eval_batches,
-                seed=args.seed + 123_456,
-            )
-
-            if is_main_process():
-                train_ppl = math.exp(
-                    min(train_loss, 20.0)
-                )
-
-                val_ppl = math.exp(
-                    min(val_loss, 20.0)
-                )
-
-                logger.log(
-                    f"step {completed_step:07d}: "
-                    f"train loss {train_loss:.4f}, "
-                    f"val loss {val_loss:.4f}, "
-                    f"Train PPL {train_ppl:.3f}, "
-                    f"Val PPL {val_ppl:.3f}, "
-                    f"tokens_seen={human_tokens(tokens_seen)}, "
-                    f"val_tokens={human_tokens(val_tokens)}, "
-                    f"toks/s={tokens_per_second:,.0f}, "
-                    f"{timestamp()}"
-                )
-
-                should_log_input_repr = (
-                    completed_step == 1
-                    or completed_step % args.milestone_save_interval_steps == 0
-                    or final_step
-                    or token_budget_reached
-                )
-                
-                if should_log_input_repr:
-                    log_input_representation_diagnostics(
-                        raw_model=raw_model,
-                        tokenizer=tokenizer,
-                        logger=logger,
-                        device=device,
-                        step=completed_step,
-                    )
-
-                greedy_prompts = prompts
-                sample_prompts = prompts
-                
-                greedy_generations = generate_diagnostics(
-                    raw_model=raw_model,
-                    tokenizer=tokenizer,
-                    prompts=greedy_prompts,
-                    device=device,
-                    max_new_tokens=args.generation_tokens,
-                    temperature=0.0,
-                    top_k=None,
-                )
-                
-                sample_generations = generate_diagnostics(
-                    raw_model=raw_model,
-                    tokenizer=tokenizer,
-                    prompts=sample_prompts,
-                    device=device,
-                    max_new_tokens=args.generation_tokens,
-                    temperature=0.8,
-                    top_k=50,
-                )
-                
-                for prompt, text in zip(greedy_prompts, greedy_generations):
-                    logger.log(
-                        f"step {completed_step:07d}: "
-                        f"[greedy] prompt={prompt!r} => {text}"
-                    )
-                
-                for prompt, text in zip(sample_prompts, sample_generations):
-                    logger.log(
-                        f"step {completed_step:07d}: "
-                        f"[sample t=0.8 k=50] prompt={prompt!r} => {text}"
-                    )
-
-                logger.log(
-                    f"step {completed_step:07d}: "
-                    f"LR: {lr:.8f}, "
-                    f"opt_step: {completed_step}, "
-                    f"{timestamp()}"
-                )
-                
-            if distributed:
-                dist.barrier() 
-                
-            last_diagnostic_time = now
-            interval_start_time = now
-            interval_start_tokens = tokens_seen
-            interval_loss_sum = 0.0
-            interval_loss_tokens = 0
-
-        save_latest = (
-            completed_step % args.latest_save_interval_steps == 0
-            or final_step
-            or token_budget_reached
-        )
-        
-        save_milestone = (
-            completed_step % args.milestone_save_interval_steps == 0
-            or final_step
-            or token_budget_reached
-        )
-
-        if save_latest or save_milestone:
-            if distributed:
-                dist.barrier()
-        
-        if save_latest:
-            save_checkpoint(
-                path=os.path.join(
-                    args.output_dir,
-                    "checkpoint_latest.pt",
-                ),
-                raw_model=raw_model,
-                optimizer=optimizer,
-                step=completed_step,
-                tokens_seen=tokens_seen,
-                args=args,
-            )
-            
-            save_model_safetensors(
-                path=os.path.join(
-                    args.output_dir,
-                    "model_latest.safetensors",
-                ),
-                raw_model=raw_model,
-            )
-        
-        if save_milestone:
-            save_checkpoint(
-                path=os.path.join(
-                    args.output_dir,
-                    f"checkpoint_{completed_step:07d}.pt",
-                ),
-                raw_model=raw_model,
-                optimizer=optimizer,
-                step=completed_step,
-                tokens_seen=tokens_seen,
-                args=args,
-            )
-            
-            save_model_safetensors(
-                path=os.path.join(
-                    args.output_dir,
-                    f"model_{completed_step:07d}.safetensors",
-                ),
-                raw_model=raw_model,
-            )
-
-        if save_latest or save_milestone:
-            if distributed:
-                dist.barrier()
-
-        if token_budget_reached:
-            if is_main_process():
-                logger.log(
-                    "[stop] max_tokens reached: "
-                    f"{tokens_seen:,}"
-                )
-            break
-
-    cleanup_distributed()
-
-
-if __name__ == "__main__":
-    main()

training_original/classic_model.py

@@ -1,594 +0,0 @@
-#!/usr/bin/env python3
-
-import math
-from typing import Optional
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from transformers import PretrainedConfig, PreTrainedModel
-from transformers.modeling_outputs import CausalLMOutput
-
-
-class ClassicConfig(PretrainedConfig):
-    model_type = "classic_causal_lm"
-
-    def __init__(
-        self,
-        vocab_size: int = 49152,
-        d_model: int = 960,
-        n_layer: int = 24,
-        n_head: int = 15,
-        ffn_multiplier: float = 8.0 / 3.0,
-        multiple_of: int = 256,
-        block_size: int = 2048,
-        rope_theta: float = 10000.0,
-        dropout: float = 0.0,
-        rms_norm_eps: float = 1e-5,
-        initializer_range: float = 0.02,
-        pad_token_id: Optional[int] = None,
-        bos_token_id: Optional[int] = None,
-        eos_token_id: Optional[int] = None,
-        tie_word_embeddings: bool = False,
-        attention_bias: bool = False,
-        mlp_bias: bool = False,
-        **kwargs,
-    ):
-        super().__init__(
-            pad_token_id=pad_token_id,
-            bos_token_id=bos_token_id,
-            eos_token_id=eos_token_id,
-            tie_word_embeddings=tie_word_embeddings,
-            **kwargs,
-        )
-
-        if d_model % n_head != 0:
-            raise ValueError("d_model must be divisible by n_head")
-
-        head_dim = d_model // n_head
-        if head_dim % 2 != 0:
-            raise ValueError("RoPE requires an even head_dim")
-
-        self.vocab_size = vocab_size
-        self.d_model = d_model
-        self.hidden_size = d_model
-        self.n_layer = n_layer
-        self.num_hidden_layers = n_layer
-        self.n_head = n_head
-        self.num_attention_heads = n_head
-        self.head_dim = head_dim
-
-        self.ffn_multiplier = ffn_multiplier
-        self.multiple_of = multiple_of
-        self.block_size = block_size
-        self.max_position_embeddings = block_size
-        self.rope_theta = rope_theta
-
-        self.dropout = dropout
-        self.rms_norm_eps = rms_norm_eps
-        self.initializer_range = initializer_range
-        self.attention_bias = attention_bias
-        self.mlp_bias = mlp_bias
-        self.use_cache = False
-
-
-class RMSNorm(nn.Module):
-    def __init__(self, dim: int, eps: float):
-        super().__init__()
-        self.weight = nn.Parameter(torch.ones(dim))
-        self.eps = eps
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        dtype = x.dtype
-        x_float = x.float()
-        x_norm = x_float * torch.rsqrt(
-            x_float.pow(2).mean(dim=-1, keepdim=True) + self.eps
-        )
-        return (x_norm * self.weight.float()).to(dtype)
-
-
-def rotate_half(x: torch.Tensor) -> torch.Tensor:
-    x1 = x[..., ::2]
-    x2 = x[..., 1::2]
-    return torch.stack((-x2, x1), dim=-1).flatten(-2)
-
-
-class RotaryEmbedding(nn.Module):
-    def __init__(self, dim: int, max_position: int, theta: float):
-        super().__init__()
-
-        inv_freq = 1.0 / (
-            theta ** (
-                torch.arange(0, dim, 2, dtype=torch.float32) / dim
-            )
-        )
-        positions = torch.arange(max_position, dtype=torch.float32)
-        freqs = torch.outer(positions, inv_freq)
-
-        # Повторяем каждую частоту для пары even/odd.
-        emb = torch.repeat_interleave(freqs, repeats=2, dim=-1)
-
-        self.register_buffer(
-            "cos_cached",
-            emb.cos(),
-            persistent=False,
-        )
-        self.register_buffer(
-            "sin_cached",
-            emb.sin(),
-            persistent=False,
-        )
-
-    def forward(
-        self,
-        q: torch.Tensor,
-        k: torch.Tensor,
-        position_ids: Optional[torch.Tensor] = None,
-    ):
-        # q, k: [B, H, T, Dh]
-        T = q.shape[-2]
-
-        if position_ids is None:
-            cos = self.cos_cached[:T][None, None, :, :]
-            sin = self.sin_cached[:T][None, None, :, :]
-        else:
-            # position_ids: [B, T]
-            cos = self.cos_cached[position_ids][:, None, :, :]
-            sin = self.sin_cached[position_ids][:, None, :, :]
-
-        cos = cos.to(device=q.device, dtype=q.dtype)
-        sin = sin.to(device=q.device, dtype=q.dtype)
-
-        q = q * cos + rotate_half(q) * sin
-        k = k * cos + rotate_half(k) * sin
-        return q, k
-
-
-class CausalSelfAttention(nn.Module):
-    def __init__(self, config: ClassicConfig):
-        super().__init__()
-
-        self.d_model = config.d_model
-        self.n_head = config.n_head
-        self.head_dim = config.d_model // config.n_head
-        self.dropout_p = config.dropout
-
-        self.q_proj = nn.Linear(
-            config.d_model,
-            config.d_model,
-            bias=config.attention_bias,
-        )
-        self.k_proj = nn.Linear(
-            config.d_model,
-            config.d_model,
-            bias=config.attention_bias,
-        )
-        self.v_proj = nn.Linear(
-            config.d_model,
-            config.d_model,
-            bias=config.attention_bias,
-        )
-        self.o_proj = nn.Linear(
-            config.d_model,
-            config.d_model,
-            bias=config.attention_bias,
-        )
-
-        self.rope = RotaryEmbedding(
-            dim=self.head_dim,
-            max_position=config.block_size,
-            theta=config.rope_theta,
-        )
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
-        B, T, C = x.shape
-
-        q = self.q_proj(x).view(
-            B, T, self.n_head, self.head_dim
-        ).transpose(1, 2)
-
-        k = self.k_proj(x).view(
-            B, T, self.n_head, self.head_dim
-        ).transpose(1, 2)
-
-        v = self.v_proj(x).view(
-            B, T, self.n_head, self.head_dim
-        ).transpose(1, 2)
-
-        q, k = self.rope(q, k, position_ids=position_ids)
-
-        dropout_p = self.dropout_p if self.training else 0.0
-
-        if attention_mask is None or bool(attention_mask.all()):
-            # is_causal=True позволяет PyTorch выбрать Flash SDP kernel.
-            out = F.scaled_dot_product_attention(
-                q,
-                k,
-                v,
-                attn_mask=None,
-                dropout_p=dropout_p,
-                is_causal=True,
-            )
-        else:
-            if attention_mask.shape != (B, T):
-                raise ValueError(
-                    f"attention_mask must be {(B, T)}, "
-                    f"got {tuple(attention_mask.shape)}"
-                )
-
-            causal = torch.ones(
-                (T, T),
-                device=x.device,
-                dtype=torch.bool,
-            ).tril()
-
-            # В SDPA bool=True означает, что элемент разрешён.
-            allowed = (
-                causal[None, None, :, :]
-                & attention_mask[:, None, None, :].bool()
-            )
-
-            out = F.scaled_dot_product_attention(
-                q,
-                k,
-                v,
-                attn_mask=allowed,
-                dropout_p=dropout_p,
-                is_causal=False,
-            )
-
-        out = out.transpose(1, 2).contiguous().view(B, T, C)
-        return self.o_proj(out)
-
-
-def round_up(value: int, multiple: int) -> int:
-    return multiple * math.ceil(value / multiple)
-
-
-class SwiGLU(nn.Module):
-    def __init__(self, config: ClassicConfig):
-        super().__init__()
-
-        hidden_dim = round_up(
-            int(config.ffn_multiplier * config.d_model),
-            config.multiple_of,
-        )
-
-        self.hidden_dim = hidden_dim
-
-        self.gate_proj = nn.Linear(
-            config.d_model,
-            hidden_dim,
-            bias=config.mlp_bias,
-        )
-        self.up_proj = nn.Linear(
-            config.d_model,
-            hidden_dim,
-            bias=config.mlp_bias,
-        )
-        self.down_proj = nn.Linear(
-            hidden_dim,
-            config.d_model,
-            bias=config.mlp_bias,
-        )
-        self.dropout = nn.Dropout(config.dropout)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = F.silu(self.gate_proj(x)) * self.up_proj(x)
-        return self.dropout(self.down_proj(x))
-
-
-class TransformerBlock(nn.Module):
-    def __init__(self, config: ClassicConfig):
-        super().__init__()
-
-        self.input_norm = RMSNorm(
-            config.d_model,
-            eps=config.rms_norm_eps,
-        )
-        self.post_attention_norm = RMSNorm(
-            config.d_model,
-            eps=config.rms_norm_eps,
-        )
-
-        self.attention = CausalSelfAttention(config)
-        self.mlp = SwiGLU(config)
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
-        x = x + self.attention(
-            self.input_norm(x),
-            attention_mask=attention_mask,
-            position_ids=position_ids,
-        )
-        x = x + self.mlp(self.post_attention_norm(x))
-        return x
-
-
-class ClassicForCausalLM(PreTrainedModel):
-    config_class = ClassicConfig
-    main_input_name = "input_ids"
-    supports_gradient_checkpointing = True
-
-    def __init__(self, config: ClassicConfig):
-        super().__init__(config)
-
-        self.token_embeddings = nn.Embedding(
-            config.vocab_size,
-            config.d_model,
-        )
-
-        self.layers = nn.ModuleList(
-            [TransformerBlock(config) for _ in range(config.n_layer)]
-        )
-
-        self.final_norm = RMSNorm(
-            config.d_model,
-            eps=config.rms_norm_eps,
-        )
-
-        self.lm_head = nn.Linear(
-            config.d_model,
-            config.vocab_size,
-            bias=False,
-        )
-
-        self.gradient_checkpointing = False
-        self.post_init()
-
-        residual_std = (
-            config.initializer_range
-            / math.sqrt(2 * config.n_layer)
-        )
-        
-        for layer in self.layers:
-            nn.init.normal_(
-                layer.attention.o_proj.weight,
-                mean=0.0,
-                std=residual_std,
-            )
-            nn.init.normal_(
-                layer.mlp.down_proj.weight,
-                mean=0.0,
-                std=residual_std,
-            )
-
-        if config.tie_word_embeddings:
-            self.tie_weights()
-
-    def _init_weights(self, module):
-        if isinstance(module, nn.Linear):
-            nn.init.normal_(
-                module.weight,
-                mean=0.0,
-                std=self.config.initializer_range,
-            )
-            if module.bias is not None:
-                nn.init.zeros_(module.bias)
-
-        elif isinstance(module, nn.Embedding):
-            nn.init.normal_(
-                module.weight,
-                mean=0.0,
-                std=self.config.initializer_range,
-            )
-
-    def _set_gradient_checkpointing(
-        self,
-        module,
-        value=False,
-        enable=None,
-        gradient_checkpointing_func=None,
-    ):
-        if enable is not None:
-            value = enable
-        self.gradient_checkpointing = value
-
-    def get_input_embeddings(self):
-        return self.token_embeddings
-
-    def set_input_embeddings(self, value):
-        self.token_embeddings = value
-
-    def get_output_embeddings(self):
-        return self.lm_head
-
-    def set_output_embeddings(self, value):
-        self.lm_head = value
-
-    def count_parameters(self):
-        total = sum(p.numel() for p in self.parameters())
-        trainable = sum(
-            p.numel() for p in self.parameters() if p.requires_grad
-        )
-        input_params = self.token_embeddings.weight.numel()
-        output_params = self.lm_head.weight.numel()
-
-        return {
-            "total": total,
-            "trainable": trainable,
-            "input": input_params,
-            "output": output_params,
-            "body": total - input_params - output_params,
-        }
-
-    def forward(
-        self,
-        input_ids: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
-        labels: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.Tensor] = None,
-        return_dict: Optional[bool] = None,
-        output_logits: bool = True,
-        **kwargs,
-    ):
-        if input_ids is None:
-            raise ValueError("input_ids must be provided")
-
-        B, T = input_ids.shape
-
-        if T > self.config.block_size:
-            raise ValueError(
-                f"Sequence length {T} exceeds "
-                f"block_size={self.config.block_size}"
-            )
-
-        if attention_mask is not None:
-            if attention_mask.shape != (B, T):
-                raise ValueError(
-                    f"attention_mask must be {(B, T)}, "
-                    f"got {tuple(attention_mask.shape)}"
-                )
-
-        x = self.token_embeddings(input_ids)
-
-        for layer in self.layers:
-            if self.gradient_checkpointing and self.training:
-                def custom_forward(
-                    hidden_states,
-                    current_layer=layer,
-                ):
-                    return current_layer(
-                        hidden_states,
-                        attention_mask=attention_mask,
-                        position_ids=position_ids,
-                    )
-        
-                x = torch.utils.checkpoint.checkpoint(
-                    custom_forward,
-                    x,
-                    use_reentrant=False,
-                )
-            else:
-                x = layer(
-                    x,
-                    attention_mask=attention_mask,
-                    position_ids=position_ids,
-                )
-
-        x = self.final_norm(x)
-        logits = self.lm_head(x) if output_logits or labels is not None else None
-
-        loss = None
-        if labels is not None:
-            loss_labels = labels.contiguous().clone()
-            
-            if attention_mask is not None:
-                loss_labels.masked_fill_(
-                    attention_mask.eq(0),
-                    -100,
-                )
-            
-            loss = F.cross_entropy(
-                logits.float().reshape(-1, self.config.vocab_size),
-                loss_labels.reshape(-1),
-                ignore_index=-100,
-            )
-            
-            '''shift_logits = logits[:, :-1, :].contiguous()
-            shift_labels = labels[:, 1:].contiguous().clone()
-
-            if attention_mask is not None:
-                shift_labels.masked_fill_(
-                    attention_mask[:, 1:].eq(0),
-                    -100,
-                )
-
-            loss = F.cross_entropy(
-                shift_logits.float().view(-1, self.config.vocab_size),
-                shift_labels.view(-1),
-                ignore_index=-100,
-            )'''
-
-        if labels is not None and labels.shape != input_ids.shape:
-            raise ValueError(
-                f"labels shape {tuple(labels.shape)} must equal "
-                f"input_ids shape {tuple(input_ids.shape)}"
-            )
-
-        return_dict = (
-            self.config.use_return_dict
-            if return_dict is None
-            else return_dict
-        )
-
-        if not return_dict:
-            output = (logits,) if output_logits else tuple()
-            return ((loss,) + output) if loss is not None else output
-
-        return CausalLMOutput(
-            loss=loss,
-            logits=logits if output_logits else None,
-        )
-
-    @torch.no_grad()
-    def generate_simple(
-        self,
-        input_ids: torch.Tensor,
-        max_new_tokens: int,
-        temperature: float = 0.0,
-        top_k: Optional[int] = None,
-        eos_token_id: Optional[int] = None,
-    ):
-        was_training = self.training
-        self.eval()
-
-        for _ in range(max_new_tokens):
-            x = input_ids[:, -self.config.block_size:]
-
-            logits = self(
-                input_ids=x,
-                return_dict=True,
-            ).logits[:, -1, :]
-
-            if temperature <= 0:
-                next_token = torch.argmax(
-                    logits,
-                    dim=-1,
-                    keepdim=True,
-                )
-            else:
-                logits = logits / temperature
-
-                if top_k is not None:
-                    values, _ = torch.topk(
-                        logits,
-                        min(top_k, logits.size(-1)),
-                    )
-                    cutoff = values[:, [-1]]
-                    logits = logits.masked_fill(
-                        logits < cutoff,
-                        float("-inf"),
-                    )
-
-                probs = F.softmax(logits.float(), dim=-1)
-                next_token = torch.multinomial(
-                    probs,
-                    num_samples=1,
-                )
-
-            input_ids = torch.cat(
-                [input_ids, next_token],
-                dim=1,
-            )
-
-            if (
-                eos_token_id is not None
-                and bool((next_token == eos_token_id).all())
-            ):
-                break
-
-        if was_training:
-            self.train()
-
-        return input_ids

training_original/classic_train.py

@@ -1,1779 +0,0 @@
-#!/usr/bin/env python3
-
-import argparse
-import contextlib
-import glob
-import json
-import math
-import os
-import random
-import time
-import glob
-from dataclasses import asdict, dataclass
-from datetime import datetime
-from pathlib import Path
-from typing import Dict, List, Optional, Tuple
-
-import numpy as np
-import torch
-import torch.distributed as dist
-from safetensors.torch import save_file
-from torch.nn.parallel import DistributedDataParallel as DDP
-from transformers import AutoTokenizer
-
-from classic_model import ClassicConfig, ClassicForCausalLM
-
-
-# ------------------------------------------------------------
-# Distributed utilities
-# ------------------------------------------------------------
-
-def distributed_is_initialized() -> bool:
-    return dist.is_available() and dist.is_initialized()
-
-
-def get_rank() -> int:
-    return dist.get_rank() if distributed_is_initialized() else 0
-
-
-def get_world_size() -> int:
-    return dist.get_world_size() if distributed_is_initialized() else 1
-
-
-def is_main_process() -> bool:
-    return get_rank() == 0
-
-
-def setup_distributed():
-    world_size = int(os.environ.get("WORLD_SIZE", "1"))
-    distributed = world_size > 1
-
-    if distributed:
-        local_rank = int(os.environ["LOCAL_RANK"])
-        torch.cuda.set_device(local_rank)
-
-        dist.init_process_group(
-            backend="nccl",
-            init_method="env://",
-        )
-
-        device = torch.device("cuda", local_rank)
-    else:
-        local_rank = 0
-        device = torch.device(
-            "cuda" if torch.cuda.is_available() else "cpu"
-        )
-
-    return distributed, local_rank, device
-
-
-def cleanup_distributed():
-    if distributed_is_initialized():
-        dist.barrier()
-        dist.destroy_process_group()
-
-
-def all_reduce_sum(value: torch.Tensor) -> torch.Tensor:
-    if distributed_is_initialized():
-        dist.all_reduce(value, op=dist.ReduceOp.SUM)
-    return value
-
-
-# ------------------------------------------------------------
-# Logging
-# ------------------------------------------------------------
-
-class Logger:
-    def __init__(self, path: str):
-        self.path = path
-
-        if is_main_process():
-            Path(path).parent.mkdir(
-                parents=True,
-                exist_ok=True,
-            )
-
-    def log(self, message: str):
-        if not is_main_process():
-            return
-
-        print(message, flush=True)
-
-        with open(self.path, "a", encoding="utf-8") as f:
-            f.write(message.rstrip() + "\n")
-
-
-def timestamp() -> str:
-    return datetime.now().strftime("%Y-%m-%d %H:%M")
-
-
-def human_tokens(value: int) -> str:
-    if value >= 1_000_000_000:
-        return f"{value / 1_000_000_000:.3f}B"
-    if value >= 1_000_000:
-        return f"{value / 1_000_000:.3f}M"
-    if value >= 1_000:
-        return f"{value / 1_000:.3f}K"
-    return str(value)
-
-
-# ------------------------------------------------------------
-# Dataset
-# ------------------------------------------------------------
-
-class TokenShardStore:
-    """
-    Stateful sequential shard sampler для файлов:
-
-        {split}_u16_*.pt
-        {split}_doc_offsets_u64_*.pt
-
-    Каждый DDP rank:
-      1. получает собственный порядок shards;
-      2. загружает один shard;
-      3. использует его для batches_per_shard microbatches;
-      4. только затем переходит к следующему shard.
-
-    Это предотвращает постоянный torch.load() с NFS и резкий
-    дисбаланс между DDP ranks.
-
-    Samples не пересекают границы документов.
-    """
-
-    def __init__(
-        self,
-        data_dir: str,
-        split: str,
-        sequence_length: int,
-        cache_shards: int = 2,
-        seed: int = 42,
-        batches_per_shard: int = 256,
-        min_long_documents: int = 1,
-    ):
-        if split not in ("train", "valid"):
-            raise ValueError("split must be 'train' or 'valid'")
-
-        if sequence_length < 1:
-            raise ValueError("sequence_length must be positive")
-
-        if batches_per_shard < 1:
-            raise ValueError("batches_per_shard must be positive")
-
-        self.data_dir = data_dir
-        self.split = split
-        self.sequence_length = int(sequence_length)
-        self.required_tokens = self.sequence_length + 1
-        self.cache_shards = max(1, int(cache_shards))
-        self.seed = int(seed)
-        self.batches_per_shard = int(batches_per_shard)
-        self.min_long_documents = int(min_long_documents)
-
-        # Значения берутся после dist.init_process_group().
-        self.rank = get_rank()
-        self.world_size = get_world_size()
-
-        token_pattern = os.path.join(
-            data_dir,
-            f"{split}_u16_*.pt",
-        )
-
-        self.token_paths = sorted(glob.glob(token_pattern))
-
-        if not self.token_paths:
-            raise FileNotFoundError(
-                f"No token shards found: {token_pattern}"
-            )
-
-        self.offset_paths = []
-
-        for token_path in self.token_paths:
-            token_name = os.path.basename(token_path)
-            suffix = token_name[len(f"{split}_u16_"):]
-
-            offset_name = (
-                f"{split}_doc_offsets_u64_{suffix}"
-            )
-            offset_path = os.path.join(
-                data_dir,
-                offset_name,
-            )
-
-            if not os.path.exists(offset_path):
-                raise FileNotFoundError(
-                    f"Missing offsets file: {offset_path}"
-                )
-
-            self.offset_paths.append(offset_path)
-
-        self.num_shards = len(self.token_paths)
-
-        # Небольшой LRU cache. Обычно активен один shard, предыдущий
-        # может остаться для дешёвого возврата после reshuffle.
-        self._cache = {}
-        self._cache_order = []
-
-        # Stateful cursor.
-        self._epoch = 0
-        self._shard_order = []
-        self._shard_cursor = 0
-
-        self._current_shard_index = None
-        self._current_tokens = None
-        self._current_offsets = None
-        self._current_valid_docs = None
-        self._batches_on_current_shard = 0
-
-        self._build_shard_order()
-        self._advance_to_next_usable_shard()
-
-    def __len__(self):
-        return self.num_shards
-
-    def _build_shard_order(self):
-        """
-        Создаёт общий deterministic permutation, затем каждый rank
-        начинает с собственного смещения.
-
-        При world_size=2 rank 0 и rank 1 обычно читают разные shards.
-        """
-        generator = torch.Generator(device="cpu")
-        generator.manual_seed(
-            self.seed + self._epoch * 1_000_003
-        )
-
-        permutation = torch.randperm(
-            self.num_shards,
-            generator=generator,
-        ).tolist()
-
-        if self.num_shards > 1:
-            shift = self.rank % self.num_shards
-            permutation = (
-                permutation[shift:]
-                + permutation[:shift]
-            )
-
-        self._shard_order = permutation
-        self._shard_cursor = 0
-        self._epoch += 1
-
-    def _touch_cache(self, shard_index: int):
-        if shard_index in self._cache_order:
-            self._cache_order.remove(shard_index)
-
-        self._cache_order.append(shard_index)
-
-        while len(self._cache_order) > self.cache_shards:
-            old_index = self._cache_order.pop(0)
-
-            if old_index == self._current_shard_index:
-                # Активный shard не удаляем.
-                self._cache_order.append(old_index)
-                break
-
-            self._cache.pop(old_index, None)
-
-    def _load_shard(
-        self,
-        shard_index: int,
-    ):
-        if shard_index in self._cache:
-            self._touch_cache(shard_index)
-            return self._cache[shard_index]
-
-        tokens = torch.load(
-            self.token_paths[shard_index],
-            map_location="cpu",
-            weights_only=True,
-        )
-
-        offsets = torch.load(
-            self.offset_paths[shard_index],
-            map_location="cpu",
-            weights_only=True,
-        )
-
-        if tokens.ndim != 1:
-            raise RuntimeError(
-                f"Token tensor must be 1D: "
-                f"{self.token_paths[shard_index]}"
-            )
-
-        if offsets.ndim != 1:
-            raise RuntimeError(
-                f"Offsets tensor must be 1D: "
-                f"{self.offset_paths[shard_index]}"
-            )
-
-        if tokens.dtype != torch.uint16:
-            raise TypeError(
-                f"Expected uint16 tokens, got {tokens.dtype}: "
-                f"{self.token_paths[shard_index]}"
-            )
-
-        if offsets.dtype not in (
-            torch.uint64,
-            torch.int64,
-        ):
-            raise TypeError(
-                f"Expected uint64/int64 offsets, "
-                f"got {offsets.dtype}: "
-                f"{self.offset_paths[shard_index]}"
-            )
-
-        if offsets.numel() < 2:
-            raise RuntimeError(
-                f"Offsets contain no documents: "
-                f"{self.offset_paths[shard_index]}"
-            )
-
-        if int(offsets[0]) != 0:
-            raise RuntimeError(
-                f"First offset must be zero: "
-                f"{self.offset_paths[shard_index]}"
-            )
-
-        if int(offsets[-1]) != tokens.numel():
-            raise RuntimeError(
-                f"Final offset/token mismatch in shard "
-                f"{shard_index}: "
-                f"{int(offsets[-1])} != {tokens.numel()}"
-            )
-
-        # Рассчитывается один раз при загрузке shard.
-        # valid_docs содержит индексы документов, которые достаточно
-        # длинные для sequence_length + target token.
-        lengths = (
-            offsets[1:].to(torch.int64)
-            - offsets[:-1].to(torch.int64)
-        )
-
-        valid_docs = torch.nonzero(
-            lengths >= self.required_tokens,
-            as_tuple=False,
-        ).flatten()
-
-        result = (
-            tokens,
-            offsets,
-            valid_docs,
-        )
-
-        self._cache[shard_index] = result
-        self._touch_cache(shard_index)
-
-        return result
-
-    def _next_shard_index(self) -> int:
-        if self._shard_cursor >= len(self._shard_order):
-            self._build_shard_order()
-
-        shard_index = self._shard_order[
-            self._shard_cursor
-        ]
-        self._shard_cursor += 1
-
-        return int(shard_index)
-
-    def _advance_to_next_usable_shard(self):
-        """
-        Переключается на следующий shard с достаточным количеством
-        длинных документов.
-
-        Ограничение attempts предотвращает бесконечный цикл, если
-        sequence_length слишком велик для всех документов.
-        """
-        attempts = 0
-        max_attempts = max(
-            self.num_shards * 2,
-            1,
-        )
-
-        while attempts < max_attempts:
-            attempts += 1
-            shard_index = self._next_shard_index()
-
-            tokens, offsets, valid_docs = self._load_shard(
-                shard_index
-            )
-
-            if valid_docs.numel() < self.min_long_documents:
-                continue
-
-            self._current_shard_index = shard_index
-            self._current_tokens = tokens
-            self._current_offsets = offsets
-            self._current_valid_docs = valid_docs
-            self._batches_on_current_shard = 0
-            return
-
-        raise RuntimeError(
-            "Could not find a usable shard containing documents "
-            f"with at least {self.required_tokens} tokens. "
-            "Reduce sequence_length or inspect document offsets."
-        )
-
-    def _ensure_current_shard(self):
-        if self._current_tokens is None:
-            self._advance_to_next_usable_shard()
-            return
-
-        if (
-            self._batches_on_current_shard
-            >= self.batches_per_shard
-        ):
-            self._advance_to_next_usable_shard()
-
-    def sample_batch(
-        self,
-        batch_size: int,
-        generator: torch.Generator,
-    ) -> torch.Tensor:
-        """
-        Возвращает LongTensor:
-
-            [batch_size, sequence_length + 1]
-
-        Все samples данного microbatch берутся из одного resident shard,
-        но из случайных документов и случайных позиций внутри документов.
-        """
-        if batch_size < 1:
-            raise ValueError("batch_size must be positive")
-
-        self._ensure_current_shard()
-
-        tokens = self._current_tokens
-        offsets = self._current_offsets
-        valid_docs = self._current_valid_docs
-
-        num_valid_docs = valid_docs.numel()
-
-        if num_valid_docs == 0:
-            raise RuntimeError(
-                "Internal error: current shard has no valid docs"
-            )
-
-        # Сразу выбираем batch_size документов.
-        selected_positions = torch.randint(
-            low=0,
-            high=num_valid_docs,
-            size=(batch_size,),
-            generator=generator,
-        )
-
-        selected_docs = valid_docs[selected_positions]
-        samples = []
-
-        for doc_tensor in selected_docs:
-            doc_index = int(doc_tensor)
-
-            begin = int(offsets[doc_index])
-            end = int(offsets[doc_index + 1])
-            document_length = end - begin
-
-            max_local_start = (
-                document_length - self.required_tokens
-            )
-
-            if max_local_start > 0:
-                local_start = int(
-                    torch.randint(
-                        low=0,
-                        high=max_local_start + 1,
-                        size=(1,),
-                        generator=generator,
-                    ).item()
-                )
-            else:
-                local_start = 0
-
-            start = begin + local_start
-            stop = start + self.required_tokens
-
-            sample = tokens[start:stop]
-
-            if sample.numel() != self.required_tokens:
-                raise RuntimeError(
-                    "Internal slicing error: "
-                    f"expected {self.required_tokens}, "
-                    f"got {sample.numel()}"
-                )
-
-            # uint16 → int64 для embedding lookup.
-            samples.append(sample.to(torch.long))
-
-        self._batches_on_current_shard += 1
-
-        return torch.stack(
-            samples,
-            dim=0,
-        )
-
-    def state_dict(self):
-        """
-        Опциональное состояние sampler для точного resume.
-        Сейчас trainer его не сохраняет, но интерфейс оставлен.
-        """
-        return {
-            "epoch": self._epoch,
-            "shard_order": list(self._shard_order),
-            "shard_cursor": self._shard_cursor,
-            "current_shard_index":
-                self._current_shard_index,
-            "batches_on_current_shard":
-                self._batches_on_current_shard,
-        }
-
-    def load_state_dict(self, state):
-        """
-        Восстанавливает shard cursor. RNG generator trainer сохраняется
-        отдельно только если вы добавите его state в checkpoint.
-        """
-        self._epoch = int(state["epoch"])
-        self._shard_order = list(state["shard_order"])
-        self._shard_cursor = int(state["shard_cursor"])
-
-        current_shard_index = state.get(
-            "current_shard_index"
-        )
-
-        if current_shard_index is None:
-            self._current_shard_index = None
-            self._current_tokens = None
-            self._current_offsets = None
-            self._current_valid_docs = None
-            self._batches_on_current_shard = 0
-            self._advance_to_next_usable_shard()
-            return
-
-        tokens, offsets, valid_docs = self._load_shard(
-            int(current_shard_index)
-        )
-
-        if valid_docs.numel() == 0:
-            raise RuntimeError(
-                "Saved current shard is no longer usable"
-            )
-
-        self._current_shard_index = int(
-            current_shard_index
-        )
-        self._current_tokens = tokens
-        self._current_offsets = offsets
-        self._current_valid_docs = valid_docs
-        self._batches_on_current_shard = int(
-            state.get(
-                "batches_on_current_shard",
-                0,
-            )
-        )
-
-
-# ------------------------------------------------------------
-# Schedules and optimizer
-# ------------------------------------------------------------
-
-def get_learning_rate(
-    step: int,
-    max_steps: int,
-    warmup_steps: int,
-    learning_rate: float,
-    min_learning_rate: float,
-) -> float:
-    if step < warmup_steps:
-        return learning_rate * float(step + 1) / max(1, warmup_steps)
-
-    if step >= max_steps:
-        return min_learning_rate
-
-    decay_ratio = (
-        step - warmup_steps
-    ) / max(1, max_steps - warmup_steps)
-
-    coefficient = 0.5 * (
-        1.0 + math.cos(math.pi * decay_ratio)
-    )
-
-    return (
-        min_learning_rate
-        + coefficient
-        * (learning_rate - min_learning_rate)
-    )
-
-
-def configure_optimizer(
-    model: torch.nn.Module,
-    learning_rate: float,
-    weight_decay: float,
-    betas: Tuple[float, float],
-    fused: bool,
-):
-    decay_params = []
-    no_decay_params = []
-
-    for name, parameter in model.named_parameters():
-        if not parameter.requires_grad:
-            continue
-
-        if parameter.dim() >= 2:
-            decay_params.append(parameter)
-        else:
-            no_decay_params.append(parameter)
-
-    groups = [
-        {
-            "params": decay_params,
-            "weight_decay": weight_decay,
-        },
-        {
-            "params": no_decay_params,
-            "weight_decay": 0.0,
-        },
-    ]
-
-    kwargs = dict(
-        lr=learning_rate,
-        betas=betas,
-        eps=1e-8,
-    )
-
-    if fused and torch.cuda.is_available():
-        kwargs["fused"] = True
-
-    return torch.optim.AdamW(groups, **kwargs)
-
-
-# ------------------------------------------------------------
-# Evaluation and generation
-# ------------------------------------------------------------
-
-@torch.no_grad()
-def estimate_loss(
-    model,
-    dataset: TokenShardStore,
-    device: torch.device,
-    batch_size: int,
-    eval_batches: int,
-    seed: int,
-) -> Tuple[float, int]:
-    model.eval()
-
-    generator = torch.Generator(device="cpu")
-    generator.manual_seed(seed + get_rank())
-
-    loss_sum = torch.zeros(
-        1,
-        device=device,
-        dtype=torch.float64,
-    )
-    token_count = torch.zeros(
-        1,
-        device=device,
-        dtype=torch.float64,
-    )
-
-    for _ in range(eval_batches):
-        batch = dataset.sample_batch(
-            batch_size=batch_size,
-            generator=generator,
-        )
-
-        batch = batch.to(
-            device,
-            non_blocking=True,
-        )
-
-        #inputs = batch[:, :-1]
-        #labels = inputs
-        
-        inputs = batch[:, :-1]
-        labels = batch[:, 1:]
-
-        with torch.autocast(
-            device_type="cuda",
-            dtype=torch.bfloat16,
-            enabled=device.type == "cuda",
-        ):
-            outputs = model(
-                input_ids=inputs,
-                labels=labels,
-                return_dict=True,
-            )
-
-        #count = labels[:, 1:].numel()
-
-        count = labels.numel()
-
-        loss_sum += outputs.loss.double() * count
-        token_count += count
-
-    all_reduce_sum(loss_sum)
-    all_reduce_sum(token_count)
-
-    mean_loss = float((loss_sum / token_count).item())
-    total_tokens = int(token_count.item())
-
-    model.train()
-    return mean_loss, total_tokens
-
-
-@torch.no_grad()
-def generate_diagnostics(
-    raw_model,
-    tokenizer,
-    prompts,
-    device,
-    max_new_tokens,
-    temperature=0.0,
-    top_k=None,
-):
-    raw_model.eval()
-    outputs = []
-
-    for prompt in prompts:
-        encoded = tokenizer(
-            prompt,
-            return_tensors="pt",
-            add_special_tokens=False,
-        )
-
-        input_ids = encoded["input_ids"].to(device)
-
-        generated = raw_model.generate_simple(
-            input_ids=input_ids,
-            max_new_tokens=max_new_tokens,
-            temperature=temperature,
-            top_k=top_k,
-            eos_token_id=tokenizer.eos_token_id,
-        )
-
-        text = tokenizer.decode(
-            generated[0],
-            skip_special_tokens=True,
-        )
-
-        outputs.append(text)
-
-    raw_model.train()
-    return outputs
-
-
-@torch.no_grad()
-def log_input_representation_diagnostics(
-    raw_model,
-    tokenizer,
-    logger,
-    device: torch.device,
-    step: int,
-    texts=None,
-    vector_prefix: int = 16,
-):
-    if not is_main_process():
-        return
-
-    if texts is None:
-        texts = [
-            "A",
-            " A",
-            "А",   # Кириллическая A
-            " А",
-            "0",
-            " 0",
-            "1",
-            " 1",
-            "the",
-            " the",
-        ]
-
-    embedding = raw_model.get_input_embeddings()
-
-    embedding_parameters = list(embedding.parameters())
-    trainable_input_parameters = sum(
-        p.numel()
-        for p in embedding_parameters
-        if p.requires_grad
-    )
-
-    logger.log(
-        f"step {step:07d}: [input_repr] "
-        f"module={embedding.__class__.__name__}, "
-        f"trainable_input_parameters={trainable_input_parameters:,}"
-    )
-
-    for text in texts:
-        token_ids = tokenizer.encode(
-            text,
-            add_special_tokens=False,
-        )
-
-        if not token_ids:
-            logger.log(
-                f"step {step:07d}: [input_repr] "
-                f"text={text!r}, token_ids=[]"
-            )
-            continue
-
-        ids = torch.tensor(
-            [token_ids],
-            dtype=torch.long,
-            device=device,
-        )
-
-        vectors = embedding(ids)[0].detach().float().cpu()
-
-        pieces = tokenizer.convert_ids_to_tokens(token_ids)
-
-        for local_index, token_id in enumerate(token_ids):
-            vector = vectors[local_index]
-            prefix = vector[:vector_prefix].tolist()
-
-            logger.log(
-                f"step {step:07d}: [input_repr] "
-                f"text={text!r}, "
-                f"piece_index={local_index}, "
-                f"token_id={token_id}, "
-                f"token_piece={pieces[local_index]!r}, "
-                f"first_{vector_prefix}="
-                f"{[round(x, 6) for x in prefix]}, "
-                f"mean={vector.mean().item():.6f}, "
-                f"std={vector.std(unbiased=False).item():.6f}, "
-                f"norm={vector.norm().item():.6f}"
-            )
-
-
-# ------------------------------------------------------------
-# Checkpoints
-# ------------------------------------------------------------
-
-def save_model_safetensors(path, raw_model):
-    if not is_main_process():
-        return
-
-    state = {
-        name: tensor.detach().cpu().contiguous()
-        for name, tensor in raw_model.state_dict().items()
-    }
-
-    tmp_path = path + ".tmp"
-    save_file(state, tmp_path)
-    os.replace(tmp_path, path)
-
-
-def save_checkpoint(
-    path: str,
-    raw_model: ClassicForCausalLM,
-    optimizer,
-    step: int,
-    tokens_seen: int,
-    args,
-):
-    if not is_main_process():
-        return
-
-    Path(path).parent.mkdir(parents=True, exist_ok=True)
-
-    tmp_path = path + ".tmp"
-
-    checkpoint = {
-        "step": step,
-        "tokens_seen": tokens_seen,
-        "model": raw_model.state_dict(),
-        "optimizer": optimizer.state_dict(),
-        "config": raw_model.config.to_dict(),
-        "args": vars(args),
-        "torch_rng_state": torch.get_rng_state(),
-        "cuda_rng_state": (
-            torch.cuda.get_rng_state_all()
-            if torch.cuda.is_available()
-            else None
-        ),
-    }
-
-    torch.save(checkpoint, tmp_path)
-    os.replace(tmp_path, path)
-
-
-def load_checkpoint(
-    path: str,
-    raw_model: ClassicForCausalLM,
-    optimizer,
-    device: torch.device,
-):
-    checkpoint = torch.load(
-        path,
-        map_location=device,
-        weights_only=False,
-    )
-
-    raw_model.load_state_dict(checkpoint["model"])
-    optimizer.load_state_dict(checkpoint["optimizer"])
-
-    if "torch_rng_state" in checkpoint:
-        torch.set_rng_state(checkpoint["torch_rng_state"].cpu())
-
-    if (
-        torch.cuda.is_available()
-        and checkpoint.get("cuda_rng_state") is not None
-    ):
-        torch.cuda.set_rng_state_all(
-            checkpoint["cuda_rng_state"]
-        )
-
-    return (
-        int(checkpoint.get("step", 0)),
-        int(checkpoint.get("tokens_seen", 0)),
-    )
-
-
-# ------------------------------------------------------------
-# CLI
-# ------------------------------------------------------------
-
-def parse_args():
-    parser = argparse.ArgumentParser()
-
-    # Paths
-    parser.add_argument("--data_dir", required=True)
-    parser.add_argument("--output_dir", required=True)
-    parser.add_argument(
-        "--tokenizer",
-        default="HuggingFaceTB/SmolLM2-135M",
-    )
-    parser.add_argument("--tokenizer_revision", default=None)
-    parser.add_argument("--resume", default=None)
-
-    # Architecture
-    parser.add_argument("--d_model", type=int, default=960)
-    parser.add_argument("--n_layer", type=int, default=24)
-    parser.add_argument("--n_head", type=int, default=15)
-    parser.add_argument(
-        "--ffn_multiplier",
-        type=float,
-        default=8.0 / 3.0,
-    )
-    parser.add_argument("--multiple_of", type=int, default=256)
-    parser.add_argument("--sequence_length", type=int, default=2048)
-    parser.add_argument("--rope_theta", type=float, default=10000.0)
-    parser.add_argument("--dropout", type=float, default=0.0)
-    parser.add_argument("--rms_norm_eps", type=float, default=1e-5)
-    parser.add_argument(
-        "--tie_word_embeddings",
-        action="store_true",
-    )
-    parser.add_argument(
-        "--gradient_checkpointing",
-        action="store_true",
-    )
-    parser.add_argument(
-        "--compile",
-        action="store_true",
-    )
-
-    # Training
-    parser.add_argument(
-        "--micro_batch_size",
-        type=int,
-        default=2,
-    )
-    parser.add_argument(
-        "--gradient_accumulation_steps",
-        type=int,
-        default=16,
-    )
-    parser.add_argument("--max_steps", type=int, default=200_000)
-    parser.add_argument(
-        "--max_tokens",
-        type=int,
-        default=0,
-        help="0 disables token-based stopping",
-    )
-    parser.add_argument(
-        "--learning_rate",
-        type=float,
-        default=3e-4,
-    )
-    parser.add_argument(
-        "--min_learning_rate",
-        type=float,
-        default=3e-5,
-    )
-    parser.add_argument("--warmup_steps", type=int, default=2000)
-    parser.add_argument("--weight_decay", type=float, default=0.1)
-    parser.add_argument("--beta1", type=float, default=0.9)
-    parser.add_argument("--beta2", type=float, default=0.95)
-    parser.add_argument("--grad_clip", type=float, default=1.0)
-    parser.add_argument("--seed", type=int, default=42)
-
-    # Data memory
-    parser.add_argument(
-        "--train_cache_shards",
-        type=int,
-        default=2,
-    )
-    parser.add_argument(
-        "--valid_cache_shards",
-        type=int,
-        default=2,
-    )
-
-    # Diagnostics
-    parser.add_argument(
-        "--diagnostic_interval_seconds",
-        type=float,
-        default=3600.0,
-    )
-    parser.add_argument(
-        "--log_interval_steps",
-        type=int,
-        default=20,
-    )
-    parser.add_argument(
-        "--eval_batches",
-        type=int,
-        default=32,
-    )
-    parser.add_argument(
-        "--eval_batch_size",
-        type=int,
-        default=2,
-    )
-    parser.add_argument(
-        "--generation_tokens",
-        type=int,
-        default=24,
-    )
-    parser.add_argument(
-        "--save_interval_steps",
-        type=int,
-        default=2000,
-    )
-    parser.add_argument(
-        "--batches_per_shard",
-        type=int,
-        default=256,
-        help=(
-            "Number of microbatches sampled from one resident shard "
-            "before loading the next shard."
-        ),
-    )
-    parser.add_argument(
-        "--latest_save_interval_steps",
-        type=int,
-        default=2000,
-    )
-    
-    parser.add_argument(
-        "--milestone_save_interval_steps",
-        type=int,
-        default=20000,
-    )
-    
-    return parser.parse_args()
-
-
-def find_nonfinite_gradients(model, max_names=20):
-    bad = []
-
-    for name, parameter in model.named_parameters():
-        gradient = parameter.grad
-
-        if gradient is None:
-            continue
-
-        finite = torch.isfinite(gradient)
-
-        if not bool(finite.all()):
-            bad.append(
-                {
-                    "name": name,
-                    "shape": tuple(gradient.shape),
-                    "dtype": str(gradient.dtype),
-                    "nan": int(torch.isnan(gradient).sum().item()),
-                    "inf": int(torch.isinf(gradient).sum().item()),
-                }
-            )
-
-            if len(bad) >= max_names:
-                break
-
-    return bad
-
-
-def grad_norm_for_named_parameters(named_parameters):
-    squares = []
-
-    for _, parameter in named_parameters:
-        if parameter.grad is None:
-            continue
-
-        gradient = parameter.grad.detach().float()
-        squares.append(gradient.pow(2).sum())
-
-    if not squares:
-        return 0.0
-
-    return torch.sqrt(torch.stack(squares).sum()).item()
-    
-
-# ------------------------------------------------------------
-# Main
-# ------------------------------------------------------------
-
-def main():
-    args = parse_args()
-
-    distributed, local_rank, device = setup_distributed()
-    rank = get_rank()
-    world_size = get_world_size()
-
-    if device.type != "cuda":
-        raise RuntimeError(
-            "This training script is intended for CUDA GPUs."
-        )
-
-    torch.backends.cuda.matmul.allow_tf32 = True
-    torch.backends.cudnn.allow_tf32 = True
-
-    # Prefer Flash SDP where available.
-    torch.backends.cuda.enable_flash_sdp(True)
-    torch.backends.cuda.enable_mem_efficient_sdp(True)
-    torch.backends.cuda.enable_math_sdp(True)
-
-    seed = args.seed + rank
-    random.seed(seed)
-    np.random.seed(seed)
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
-
-    Path(args.output_dir).mkdir(
-        parents=True,
-        exist_ok=True,
-    )
-
-    logger = Logger(
-        os.path.join(args.output_dir, "train.log")
-    )
-
-    tokenizer = AutoTokenizer.from_pretrained(
-        args.tokenizer,
-        revision=args.tokenizer_revision,
-        use_fast=True,
-    )
-
-    vocab_size = len(tokenizer)
-
-    # Не используем произвольный legacy pad ID.
-    # Dataset состоит из fixed-length samples и padding не нужен.
-    config = ClassicConfig(
-        vocab_size=vocab_size,
-        d_model=args.d_model,
-        n_layer=args.n_layer,
-        n_head=args.n_head,
-        ffn_multiplier=args.ffn_multiplier,
-        multiple_of=args.multiple_of,
-        block_size=args.sequence_length,
-        rope_theta=args.rope_theta,
-        dropout=args.dropout,
-        rms_norm_eps=args.rms_norm_eps,
-        pad_token_id=tokenizer.pad_token_id,
-        bos_token_id=tokenizer.bos_token_id,
-        eos_token_id=tokenizer.eos_token_id,
-        tie_word_embeddings=args.tie_word_embeddings,
-    )
-
-    raw_model = ClassicForCausalLM(config)
-
-    if args.gradient_checkpointing:
-        raw_model.gradient_checkpointing = True
-
-    raw_model.to(device)
-
-    if is_main_process():
-        first_parameter = next(raw_model.parameters())
-    
-        logger.log(
-            "[dtype] "
-            f"parameter_dtype={first_parameter.dtype}, "
-            f"optimizer_master_expected=float32"
-        )
-    
-    if next(raw_model.parameters()).dtype != torch.float32:
-        raise RuntimeError(
-            "Model parameters must remain FP32; "
-            "BF16 should be enabled only through autocast"
-        )
-
-    parameter_counts = raw_model.count_parameters()
-
-    if is_main_process():
-        logger.log(
-            "[model] "
-            + ", ".join(
-                f"{key}={value:,}"
-                for key, value in parameter_counts.items()
-            )
-        )
-        logger.log(
-            "[config] "
-            + json.dumps(
-                config.to_dict(),
-                ensure_ascii=False,
-                sort_keys=True,
-            )
-        )
-        logger.log(
-            f"[run] world_size={world_size}, "
-            f"micro_batch={args.micro_batch_size}, "
-            f"grad_accum={args.gradient_accumulation_steps}, "
-            f"seq={args.sequence_length}, "
-            f"global_tokens_per_step="
-            #f"{world_size * args.micro_batch_size * args.gradient_accumulation_steps * (args.sequence_length - 1):,}"
-            f"{world_size * args.micro_batch_size * args.gradient_accumulation_steps * args.sequence_length:,}"
-        )
-
-    optimizer = configure_optimizer(
-        raw_model,
-        learning_rate=args.learning_rate,
-        weight_decay=args.weight_decay,
-        betas=(args.beta1, args.beta2),
-        fused=True,
-    )
-
-    start_step = 0
-    tokens_seen = 0
-
-    if args.resume is not None:
-        start_step, tokens_seen = load_checkpoint(
-            args.resume,
-            raw_model,
-            optimizer,
-            device,
-        )
-
-        logger.log(
-            f"[resume] path={args.resume}, "
-            f"step={start_step}, "
-            f"tokens_seen={tokens_seen:,}"
-        )
-
-    model = raw_model
-
-    if args.compile:
-        model = torch.compile(
-            model,
-            mode="max-autotune",
-            dynamic=False,
-        )
-
-    if distributed:
-        '''model = DDP(
-            model,
-            device_ids=[local_rank],
-            output_device=local_rank,
-            broadcast_buffers=False,
-            gradient_as_bucket_view=True,
-            static_graph=not args.gradient_checkpointing,
-        )'''
-        model = DDP(
-            model,
-            device_ids=[local_rank],
-            output_device=local_rank,
-            broadcast_buffers=False,
-            gradient_as_bucket_view=True,
-            static_graph=False,
-            find_unused_parameters=False,
-        )
-    
-    train_data = TokenShardStore(
-        data_dir=args.data_dir,
-        split="train",
-        sequence_length=args.sequence_length,
-        cache_shards=args.train_cache_shards,
-        seed=args.seed,
-        batches_per_shard=args.batches_per_shard,
-    )
-
-    valid_data = TokenShardStore(
-        data_dir=args.data_dir,
-        split="valid",
-        sequence_length=args.sequence_length,
-        cache_shards=args.valid_cache_shards,
-        seed=args.seed + 10_000,
-        batches_per_shard=max(
-            args.batches_per_shard,
-            args.eval_batches,
-        ),
-    )
-
-    train_generator = torch.Generator(device="cpu")
-    train_generator.manual_seed(args.seed + rank * 100_003)
-
-    prompts = [
-        # English factual completion
-        "London is the capital of",
-        "The capital of France is",
-        "The largest planet in the Solar System is",
-        "Water freezes at",
-        "The chemical symbol for gold is",
-        "The Pacific Ocean is",
-        "The human heart pumps",
-        "The Second World War ended in",
-        "The author of Romeo and Juliet was",
-        "A triangle has",
-    
-        # English continuation and grammar
-        "Once upon a time, there was",
-        "The scientist opened the laboratory door and",
-        "When the rain finally stopped,",
-        "She went to the store because",
-        "If I had known about the problem,",
-        "The old house on the hill",
-        "Although the experiment failed,",
-        "In order to solve this problem, we need to",
-        "The main difference between cats and dogs is",
-        "This article explains how to",
-    
-        # Definitions and explanations
-        "Photosynthesis is the process by which",
-        "Gravity is a force that",
-        "A computer program is",
-        "Democracy can be defined as",
-        "Machine learning is used to",
-        "The purpose of a database is to",
-        "An ecosystem consists of",
-        "Inflation occurs when",
-        "The Internet allows people to",
-        "Energy cannot be created or destroyed, but",
-    
-        # Arithmetic and symbolic patterns
-        "2 + 2 =",
-        "10 - 3 =",
-        "6 * 7 =",
-        "12 / 4 =",
-        "1, 2, 3, 4,",
-        "2, 4, 6, 8,",
-        "The square root of 9 is",
-        "If x = 5, then x + 2 =",
-        "One hundred divided by ten equals",
-        "The next number after 99 is",
-    ]
-
-    # Running training loss since the previous diagnostic.
-    interval_loss_sum = 0.0
-    interval_loss_tokens = 0
-    interval_start_tokens = tokens_seen
-    interval_start_time = time.monotonic()
-
-    last_diagnostic_time = time.monotonic()
-
-    # Initial diagnostic.
-    if is_main_process():
-        logger.log(f"[start] {timestamp()}")
-        
-        log_input_representation_diagnostics(
-            raw_model=raw_model,
-            tokenizer=tokenizer,
-            logger=logger,
-            device=device,
-            step=start_step,
-        )
-
-    model.train()
-    optimizer.zero_grad(set_to_none=True)
-
-    for step in range(start_step, args.max_steps):
-        lr = get_learning_rate(
-            step=step,
-            max_steps=args.max_steps,
-            warmup_steps=args.warmup_steps,
-            learning_rate=args.learning_rate,
-            min_learning_rate=args.min_learning_rate,
-        )
-
-        for group in optimizer.param_groups:
-            group["lr"] = lr
-
-        step_loss_sum = 0.0
-        step_loss_tokens = 0
-
-        for micro_step in range(
-            args.gradient_accumulation_steps
-        ):
-            batch = train_data.sample_batch(
-                batch_size=args.micro_batch_size,
-                generator=train_generator,
-            )
-
-            batch = batch.to(
-                device,
-                non_blocking=True,
-            )
-
-            # forward() сам сдвигает logits и labels на один токен.
-            #input_ids = batch[:, :-1]
-            #labels = input_ids
-
-            input_ids = batch[:, :-1]
-            labels = batch[:, 1:]
-
-            should_sync = (
-                micro_step
-                == args.gradient_accumulation_steps - 1
-            )
-
-            sync_context = contextlib.nullcontext()
-
-            if distributed and not should_sync:
-                sync_context = model.no_sync()
-
-            with sync_context:
-                with torch.autocast(
-                    device_type="cuda",
-                    dtype=torch.bfloat16,
-                ):
-                    outputs = model(
-                        input_ids=input_ids,
-                        labels=labels,
-                        return_dict=True,
-                    )
-
-                    loss = (
-                        outputs.loss
-                        / args.gradient_accumulation_steps
-                    )
-
-                loss.backward()
-
-            # Первый label не имеет предшествующего logit после внутреннего shift.
-            
-            #local_tokens = labels[:, 1:].numel()
-            
-            local_tokens = labels.numel()
-
-            step_loss_sum += (
-                float(outputs.loss.detach()) * local_tokens
-            )
-            step_loss_tokens += local_tokens
-
-        
-        bad_gradients = find_nonfinite_gradients(raw_model)
-        
-        local_bad = torch.tensor(
-            [1 if bad_gradients else 0],
-            device=device,
-            dtype=torch.int32,
-        )
-        
-        if distributed:
-            dist.all_reduce(
-                local_bad,
-                op=dist.ReduceOp.MAX,
-            )
-        
-        if int(local_bad.item()) != 0:
-            if bad_gradients:
-                logger.log(
-                    "[nonfinite_gradients] "
-                    + json.dumps(
-                        bad_gradients,
-                        ensure_ascii=False,
-                    )
-                )
-        
-            emergency_path = os.path.join(
-                args.output_dir,
-                f"checkpoint_nonfinite_step_{step + 1:07d}.pt",
-            )
-        
-            save_checkpoint(
-                path=emergency_path,
-                raw_model=raw_model,
-                optimizer=optimizer,
-                step=step,
-                tokens_seen=tokens_seen,
-                args=args,
-            )
-        
-            raise RuntimeError(
-                f"Non-finite gradients at step {step + 1}"
-            )
-
-        input_norm = grad_norm_for_named_parameters(
-            raw_model.token_embeddings.named_parameters()
-        )
-        
-        body_norm = grad_norm_for_named_parameters(
-            (
-                (name, parameter)
-                for name, parameter in raw_model.named_parameters()
-                if not name.startswith("token_embeddings.")
-                and not name.startswith("lm_head.")
-            )
-        )
-        
-        output_norm = grad_norm_for_named_parameters(
-            raw_model.lm_head.named_parameters()
-        )
-        
-        #logger.log(
-        #    f"[grad_groups] input={input_norm:.3f}, "
-        #    f"body={body_norm:.3f}, output={output_norm:.3f}"
-        #)
-        
-        if args.grad_clip > 0:
-            grad_norm = torch.nn.utils.clip_grad_norm_(
-                raw_model.parameters(),
-                args.grad_clip,
-                error_if_nonfinite=True,
-            )
-        else:
-            grad_norm = torch.tensor(
-                float("nan"),
-                device=device,
-            )
-
-        optimizer.step()
-        optimizer.zero_grad(set_to_none=True)
-
-        # Global token count.
-        global_step_tokens = (
-            step_loss_tokens * world_size
-        )
-        tokens_seen += global_step_tokens
-
-        # Loss is averaged approximately across ranks here.
-        loss_stats = torch.tensor(
-            [step_loss_sum, step_loss_tokens],
-            device=device,
-            dtype=torch.float64,
-        )
-        all_reduce_sum(loss_stats)
-
-        global_loss_sum = float(loss_stats[0].item())
-        global_loss_tokens = int(loss_stats[1].item())
-
-        interval_loss_sum += global_loss_sum
-        interval_loss_tokens += global_loss_tokens
-
-        completed_step = step + 1
-
-        if (
-            completed_step % args.log_interval_steps == 0
-            and is_main_process()
-        ):
-            mean_step_loss = (
-                global_loss_sum / global_loss_tokens
-            )
-
-            logger.log(
-                f"step {completed_step:07d}: "
-                f"loss {mean_step_loss:.4f}, "
-                f"lr {lr:.8f}, "
-                f"grad_norm {float(grad_norm):.4f}, "
-                f"tokens_seen={human_tokens(tokens_seen)}, "
-                f"{timestamp()}"
-            )
-
-        now = time.monotonic()
-
-        '''diagnostic_due = (
-            now - last_diagnostic_time
-            >= args.diagnostic_interval_seconds
-        )'''
-
-        diagnostic_due = (
-            completed_step % 700 == 0
-        )
-
-        final_step = completed_step == args.max_steps
-
-        token_budget_reached = (
-            args.max_tokens > 0
-            and tokens_seen >= args.max_tokens
-        )
-
-        if diagnostic_due or final_step or token_budget_reached:
-            if distributed:
-                dist.barrier()
-        
-            elapsed = max(
-                now - interval_start_time,
-                1e-9,
-            )
-            interval_tokens = (
-                tokens_seen - interval_start_tokens
-            )
-            tokens_per_second = interval_tokens / elapsed
-
-            train_loss = (
-                interval_loss_sum
-                / max(interval_loss_tokens, 1)
-            )
-
-            val_loss, val_tokens = estimate_loss(
-                model=model,
-                dataset=valid_data,
-                device=device,
-                batch_size=args.eval_batch_size,
-                eval_batches=args.eval_batches,
-                seed=args.seed + 123_456,
-            )
-
-            if is_main_process():
-                train_ppl = math.exp(min(train_loss, 20.0))
-                val_ppl = math.exp(min(val_loss, 20.0))
-
-                logger.log(
-                    f"step {completed_step:07d}: "
-                    f"train loss {train_loss:.4f}, "
-                    f"val loss {val_loss:.4f}, "
-                    f"Train PPL {train_ppl:.3f}, "
-                    f"Val PPL {val_ppl:.3f}, "
-                    f"tokens_seen={human_tokens(tokens_seen)}, "
-                    f"val_tokens={human_tokens(val_tokens)}, "
-                    f"toks/s={tokens_per_second:,.0f}, "
-                    f"{timestamp()}"
-                )
-
-                should_log_input_repr = (
-                    completed_step == 1
-                    or completed_step % args.milestone_save_interval_steps == 0
-                    or final_step
-                    or token_budget_reached
-                )
-                
-                if should_log_input_repr:
-                    log_input_representation_diagnostics(
-                        raw_model=raw_model,
-                        tokenizer=tokenizer,
-                        logger=logger,
-                        device=device,
-                        step=completed_step,
-                    )
-
-                greedy_prompts = prompts
-                sample_prompts = prompts
-                
-                greedy_generations = generate_diagnostics(
-                    raw_model=raw_model,
-                    tokenizer=tokenizer,
-                    prompts=greedy_prompts,
-                    device=device,
-                    max_new_tokens=args.generation_tokens,
-                    temperature=0.0,
-                    top_k=None,
-                )
-                
-                sample_generations = generate_diagnostics(
-                    raw_model=raw_model,
-                    tokenizer=tokenizer,
-                    prompts=sample_prompts,
-                    device=device,
-                    max_new_tokens=args.generation_tokens,
-                    temperature=0.8,
-                    top_k=50,
-                )
-                
-                for prompt, text in zip(greedy_prompts, greedy_generations):
-                    logger.log(
-                        f"step {completed_step:07d}: "
-                        f"[greedy] prompt={prompt!r} => {text}"
-                    )
-                
-                for prompt, text in zip(sample_prompts, sample_generations):
-                    logger.log(
-                        f"step {completed_step:07d}: "
-                        f"[sample t=0.8 k=50] prompt={prompt!r} => {text}"
-                    )
-
-                logger.log(
-                    f"step {completed_step:07d}: "
-                    f"LR: {lr:.8f}, "
-                    f"opt_step: {completed_step}, "
-                    f"{timestamp()}"
-                )
-
-            if distributed:
-                dist.barrier() 
-
-            last_diagnostic_time = now
-            interval_start_time = now
-            interval_start_tokens = tokens_seen
-            interval_loss_sum = 0.0
-            interval_loss_tokens = 0
-
-        save_latest = (
-            completed_step % args.latest_save_interval_steps == 0
-            or final_step
-            or token_budget_reached
-        )
-        
-        save_milestone = (
-            completed_step % args.milestone_save_interval_steps == 0
-            or final_step
-            or token_budget_reached
-        )
-
-        if save_latest or save_milestone:
-            if distributed:
-                dist.barrier()
-        
-        if save_latest:
-            save_checkpoint(
-                path=os.path.join(
-                    args.output_dir,
-                    "checkpoint_latest.pt",
-                ),
-                raw_model=raw_model,
-                optimizer=optimizer,
-                step=completed_step,
-                tokens_seen=tokens_seen,
-                args=args,
-            )
-            
-            save_model_safetensors(
-                path=os.path.join(
-                    args.output_dir,
-                    "model_latest.safetensors",
-                ),
-                raw_model=raw_model,
-            )
-        
-        if save_milestone:
-            save_checkpoint(
-                path=os.path.join(
-                    args.output_dir,
-                    f"checkpoint_{completed_step:07d}.pt",
-                ),
-                raw_model=raw_model,
-                optimizer=optimizer,
-                step=completed_step,
-                tokens_seen=tokens_seen,
-                args=args,
-            )
-            
-            save_model_safetensors(
-                path=os.path.join(
-                    args.output_dir,
-                    f"model_{completed_step:07d}.safetensors",
-                ),
-                raw_model=raw_model,
-            )
-
-        if save_latest or save_milestone:
-            if distributed:
-                dist.barrier()
-
-        if token_budget_reached:
-            if is_main_process():
-                logger.log(
-                    f"[stop] max_tokens reached: "
-                    f"{tokens_seen:,}"
-                )
-            break
-
-    cleanup_distributed()
-
-
-if __name__ == "__main__":
-    main()

training_original/model_16_dim_bin.py

@@ -1,208 +0,0 @@
-#!/usr/bin/env python3
-
-from typing import Optional
-
-import torch
-import torch.nn as nn
-
-from classic_model import ClassicConfig, ClassicForCausalLM
-
-
-class Binary16Config(ClassicConfig):
-    model_type = "binary16_causal_lm"
-
-    def __init__(
-        self,
-        binary_dim: int = 16,
-        binary_encoding: str = "zero_one",
-        binary_scale: float = 1.0,
-        binary_permutation_seed: Optional[int] = None,
-        **kwargs,
-    ):
-        super().__init__(**kwargs)
-
-        if binary_dim != 16:
-            raise ValueError(
-                "This implementation requires binary_dim=16"
-            )
-
-        if self.vocab_size > (1 << binary_dim):
-            raise ValueError(
-                f"vocab_size={self.vocab_size} exceeds "
-                f"2**{binary_dim}"
-            )
-
-        if self.d_model % binary_dim != 0:
-            raise ValueError(
-                f"d_model={self.d_model} must be divisible "
-                f"by binary_dim={binary_dim}"
-            )
-
-        if binary_encoding not in ("zero_one", "bipolar"):
-            raise ValueError(
-                "binary_encoding must be zero_one or bipolar"
-            )
-
-        if self.tie_word_embeddings:
-            raise ValueError(
-                "tie_word_embeddings must be False for "
-                "fixed binary input"
-            )
-
-        self.binary_dim = binary_dim
-        self.binary_encoding = binary_encoding
-        self.binary_scale = float(binary_scale)
-        self.binary_permutation_seed = binary_permutation_seed
-        self.binary_repeat = self.d_model // binary_dim
-
-
-def build_binary_codebook(
-    vocab_size: int,
-    binary_dim: int,
-    encoding: str,
-    permutation_seed: Optional[int],
-) -> torch.Tensor:
-    if vocab_size > (1 << binary_dim):
-        raise ValueError(
-            f"vocab_size={vocab_size} does not fit "
-            f"in {binary_dim} bits"
-        )
-
-    if permutation_seed is None:
-        code_ids = torch.arange(
-            vocab_size,
-            dtype=torch.int64,
-        )
-    else:
-        generator = torch.Generator(device="cpu")
-        generator.manual_seed(permutation_seed)
-
-        code_ids = torch.randperm(
-            1 << binary_dim,
-            generator=generator,
-            dtype=torch.int64,
-        )[:vocab_size]
-
-    shifts = torch.arange(
-        binary_dim,
-        dtype=torch.int64,
-    )
-
-    codebook = (
-        (code_ids[:, None] >> shifts[None, :]) & 1
-    ).to(torch.float32)
-
-    if encoding == "bipolar":
-        codebook = codebook.mul(2.0).sub(1.0)
-    elif encoding != "zero_one":
-        raise ValueError(f"Unknown encoding: {encoding}")
-
-    return codebook.contiguous()
-
-
-class FixedBinary16Embedding(nn.Module):
-    def __init__(self, config: Binary16Config):
-        super().__init__()
-
-        codebook = build_binary_codebook(
-            vocab_size=config.vocab_size,
-            binary_dim=config.binary_dim,
-            encoding=config.binary_encoding,
-            permutation_seed=config.binary_permutation_seed,
-        )
-
-        # Buffer, не Parameter.
-        self.register_buffer(
-            "codebook",
-            codebook,
-            persistent=True,
-        )
-
-        self.vocab_size = config.vocab_size
-        self.binary_dim = config.binary_dim
-        self.d_model = config.d_model
-        self.repeat = config.binary_repeat
-        self.binary_scale = config.binary_scale
-
-    @property
-    def weight(self):
-        # Совместимость с get_input_embeddings().
-        # Возвращается buffer, не trainable Parameter.
-        return self.codebook
-
-    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
-        code = self.codebook[input_ids.long()]
-        x = code.repeat(
-            *([1] * (code.ndim - 1)),
-            self.repeat,
-        )
-
-        if self.binary_scale != 1.0:
-            x = x * self.binary_scale
-
-        return x
-
-class Binary16ForCausalLM(ClassicForCausalLM):
-    config_class = Binary16Config
-
-    def __init__(self, config: Binary16Config):
-        # Временно создаётся стандартная embedding-таблица,
-        # затем немедленно удаляется до optimizer construction.
-        super().__init__(config)
-
-        self.token_embeddings = FixedBinary16Embedding(
-            config
-        )
-
-    def get_input_embeddings(self):
-        return self.token_embeddings
-
-    def set_input_embeddings(self, value):
-        raise RuntimeError(
-            "Binary16ForCausalLM has a fixed input interface"
-        )
-
-    def tie_weights(self):
-        if getattr(
-            self.config,
-            "tie_word_embeddings",
-            False,
-        ):
-            raise ValueError(
-                "Fixed binary input cannot be tied "
-                "to the output projection"
-            )
-
-    def count_parameters(self):
-        total = sum(
-            parameter.numel()
-            for parameter in self.parameters()
-        )
-
-        trainable = sum(
-            parameter.numel()
-            for parameter in self.parameters()
-            if parameter.requires_grad
-        )
-
-        frozen_parameters = sum(
-            parameter.numel()
-            for parameter in self.parameters()
-            if not parameter.requires_grad
-        )
-
-        output_parameters = (
-            self.lm_head.weight.numel()
-        )
-
-        return {
-            "total_parameters": total,
-            "trainable_parameters": trainable,
-            "frozen_parameters": frozen_parameters,
-            "input_trainable_parameters": 0,
-            "fixed_codebook_values":
-                self.token_embeddings.codebook.numel(),
-            "output_parameters": output_parameters,
-            "body_parameters":
-                total - output_parameters,
-        }