Initial upload of JiRackNative 3B pre-train weights first checkpoint

JiRackTernaryPyTorch_3b.py

@@ -0,0 +1,166 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.checkpoint import checkpoint
+
+# --- JIRACK 3B CONSTANTS ---
+VOCAB_SIZE = 128256
+HIDDEN_SIZE = 3072
+NUM_LAYERS = 20
+NUM_HEADS = 24
+NUM_KV_HEADS = 8
+#INTERMEDIATE_SIZE = 8192
+INTERMEDIATE_SIZE = 4096
+MAX_SEQ_LEN = 4096
+RMS_EPS = 1e-6
+STABILITY_EPS = 1e-9
+INT8_SCALE_TARGET = 127.0
+TERNARY = False
+
+class TernaryConfig:
+    def __init__(self):
+        self.vocab_size = VOCAB_SIZE
+        self.hidden_size = HIDDEN_SIZE
+        self.num_hidden_layers = NUM_LAYERS
+        self.num_attention_heads = NUM_HEADS
+        self.num_key_value_heads = NUM_KV_HEADS
+        self.intermediate_size = INTERMEDIATE_SIZE
+        self.max_position_embeddings = MAX_SEQ_LEN
+        self.rms_norm_eps = RMS_EPS
+        self.tie_word_embeddings = False
+        self.model_type = "jirack_ternary"
+        self.ternary = TERNARY  # Флаг теперь внутри конфига
+
+    def get(self, key, default=None):
+        return getattr(self, key, default)
+
+    def __getitem__(self, key):
+        return getattr(self, key)
+
+class BitLinear(nn.Linear):
+    def __init__(self, in_features, out_features, bias=False, ternary=False):
+        super().__init__(in_features, out_features, bias)
+        self.ternary = ternary
+
+    def forward(self, x):
+        if not self.ternary:
+            return F.linear(x, self.weight, self.bias)
+        # Weight Quantization
+        w = self.weight
+        gamma = w.abs().mean().clamp(min=STABILITY_EPS)
+        w_quant = torch.clamp(torch.round(w / gamma), -1, 1)
+        w_final = w + (w_quant * gamma - w).detach()
+
+        # Activation Quantization (Absmax)
+        x_norm = x - x.mean(dim=-1, keepdim=True)
+        x_max = x_norm.abs().max(dim=-1, keepdim=True).values.clamp(min=STABILITY_EPS)
+        scale = INT8_SCALE_TARGET / x_max
+        x_quant = (x_norm * scale).round().clamp(-128, 127) / scale
+        x_final = x + (x_quant - x).detach()
+
+        return F.linear(x_final, w_final, self.bias)
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps=RMS_EPS):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+    def forward(self, x):
+        # Убран жесткий .pow(2) в float32, чтобы не ломать скорость
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight
+
+def precompute_freqs_cis(dim, seq_len, theta=500000.0):
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
+    t = torch.arange(seq_len).float()
+    freqs = torch.outer(t, freqs)
+    return torch.cos(freqs), torch.sin(freqs)
+
+def apply_rotary_emb(xq, xk, freqs_cos, freqs_sin):
+    def rotate_half(x):
+        x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :]
+        return torch.cat((-x2, x1), dim=-1)
+    T = xq.shape[2]
+    # Убрали принудительный .to(torch.float32)
+    f_cos = freqs_cos[:T].to(device=xq.device, dtype=xq.dtype).view(1, 1, T, -1).repeat(1, 1, 1, 2)
+    f_sin = freqs_sin[:T].to(device=xq.device, dtype=xq.dtype).view(1, 1, T, -1).repeat(1, 1, 1, 2)
+    return (xq * f_cos) + (rotate_half(xq) * f_sin), (xk * f_cos) + (rotate_half(xk) * f_sin)
+
+class TransformerBlock(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.n_heads = config.num_attention_heads
+        self.n_kv_heads = config.num_key_value_heads
+        self.n_rep = self.n_heads // self.n_kv_heads
+        self.head_dim = config.hidden_size // self.n_heads
+        # Передаем параметр ternary из конфигурации
+        self.q_proj = BitLinear(config.hidden_size, config.hidden_size, ternary=config.ternary)
+        self.k_proj = BitLinear(config.hidden_size, self.n_kv_heads * self.head_dim, ternary=config.ternary)
+        self.v_proj = BitLinear(config.hidden_size, self.n_kv_heads * self.head_dim, ternary=config.ternary)
+        self.out_proj = BitLinear(config.hidden_size, config.hidden_size, ternary=config.ternary)
+
+        self.ffn_w1 = BitLinear(config.hidden_size, config.intermediate_size, ternary=config.ternary)
+        self.ffn_w3 = BitLinear(config.hidden_size, config.intermediate_size, ternary=config.ternary)
+        self.ffn_w2 = BitLinear(config.intermediate_size, config.hidden_size, ternary=config.ternary)
+
+        self.norm1, self.norm2 = RMSNorm(config.hidden_size), RMSNorm(config.hidden_size)
+
+    def forward(self, x, freqs_cos, freqs_sin):
+        h = self.norm1(x)
+        B, T, D = x.shape
+
+        q = self.q_proj(h).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(h).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(h).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
+
+        q, k = apply_rotary_emb(q, k, freqs_cos, freqs_sin)
+
+        if self.n_rep > 1:
+            k = k[:, :, None, :, :].expand(B, self.n_kv_heads, self.n_rep, T, self.head_dim).reshape(B, self.n_heads, T, self.head_dim)
+            v = v[:, :, None, :, :].expand(B, self.n_kv_heads, self.n_rep, T, self.head_dim).reshape(B, self.n_heads, T, self.head_dim)
+
+        # Полностью автоматический выбор кернела силами PyTorch
+        attn_out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+
+        x = x + self.out_proj(attn_out.transpose(1, 2).reshape(B, T, D))
+        m = self.norm2(x)
+        x = x + self.ffn_w2(F.silu(self.ffn_w1(m)) * self.ffn_w3(m))
+        return x
+
+class TernaryTransformer3B(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.token_emb = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.blocks = nn.ModuleList([TransformerBlock(config) for _ in range(config.num_hidden_layers)])
+        self.ln_f = RMSNorm(config.hidden_size)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        self.head_dim = config.hidden_size // config.num_attention_heads
+        self.gradient_checkpointing = False
+        self._set_rope_cache(config.max_position_embeddings)
+        print(f"Ternary={config.ternary} | Native Auto-SDPA Activated")
+        
+    def gradient_checkpointing_enable(self, **kwargs):
+        self.gradient_checkpointing = True
+
+    def _set_rope_cache(self, seq_len):
+        cos, sin = precompute_freqs_cis(self.head_dim, seq_len)
+        self.register_buffer("freqs_cos", cos, persistent=False)
+        self.register_buffer("freqs_sin", sin, persistent=False)
+
+    def forward(self, input_ids):
+        input_ids = input_ids.to(torch.long)
+        T = input_ids.shape[1]
+        if T > self.freqs_cos.shape[0]:
+            self._set_rope_cache(T)
+
+        x = self.token_emb(input_ids)
+
+        for block in self.blocks:
+            if self.gradient_checkpointing and self.training:
+                x = checkpoint(block, x, self.freqs_cos, self.freqs_sin, use_reentrant=False)
+            else:
+                x = block(x, self.freqs_cos, self.freqs_sin)
+
+        logits = self.lm_head(self.ln_f(x))
+        return logits, None

tokenizer_config.json

@@ -0,0 +1,10 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "<|endoftext|>",
+  "clean_up_tokenization_spaces": true,
+  "eos_token": "<|endoftext|>",
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<|padding|>",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": "<|unk|>"
+}

train_jirack_accelerate.py

@@ -0,0 +1,199 @@
+import os
+# Включаем оптимизацию памяти для ROCm/HIP ДО импорта torch!
+os.environ["PYTORCH_HIP_ALLOC_CONF"] = "expandable_segments:True"
+
+import glob
+import math
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset, DataLoader
+from accelerate import Accelerator
+from tqdm import tqdm
+from transformers import Adafactor, get_cosine_schedule_with_warmup
+
+# --- ГЛОБАЛЬНЫЕ КОНСТАНТЫ ---
+USE_COSINE_SCHEDULER = True  # Переключите в False, если нужно отключить планировщик (warmup + cosine)
+
+# --- 1. Легковесный датасет ---
+class SingleShardDataset(Dataset):
+    def __init__(self, shard_path):
+        self.data = torch.load(shard_path, map_location="cpu", weights_only=True)
+    def __len__(self):
+        return self.data.shape[0]
+    def __getitem__(self, idx):
+        return self.data[idx].long()
+
+# --- 2. Основная функция тренировки ---
+def train():
+    grad_accumulation_steps = 1
+    batch_size = 1
+
+    # Инициализируем Accelerator
+    accelerator = Accelerator(
+        mixed_precision="bf16",
+        gradient_accumulation_steps=grad_accumulation_steps
+    )
+
+    pt_chunks_mask = "/mnt/nfs_clientshare/JiRackPretrain/jirack_pretrain_chunk_*.pt"
+    checkpoint_dir = "checkpoints"
+    pt_files = sorted(glob.glob(pt_chunks_mask))
+
+    if not pt_files:
+        raise FileNotFoundError(f"Не найдены файлы чанков по маске: {pt_chunks_mask}")
+
+    if accelerator.is_local_main_process:
+        print(f"Найдено шардов: {len(pt_files)}")
+        print("Инициализация JiRack 3.3B...")
+
+    # Импортируем строго ваши классы из локального файла
+    from JiRackTernaryPyTorch_3b import TernaryTransformer3B, TernaryConfig
+
+    config = TernaryConfig()
+    model = TernaryTransformer3B(config)
+    
+    # === ТОЧЕЧНАЯ ЗАГРУЗКА ЧЕКПОИНТА ===
+    checkpoint_load_path = "model_weights.pt"
+    if os.path.exists(checkpoint_load_path):
+        if accelerator.is_local_main_process:
+            print(f"-> Загрузка сохраненных весов из: {checkpoint_load_path}")
+        state_dict = torch.load(checkpoint_load_path, map_location="cpu", weights_only=True)
+        model.load_state_dict(state_dict)
+    else:
+        if accelerator.is_local_main_process:
+            print(f"-> Чекпоинт не найден по пути {checkpoint_load_path}, обучение начнется с нуля.")
+    # ==================================
+
+    model.gradient_checkpointing_enable()
+
+    if accelerator.is_local_main_process:
+        print("-> Gradient Checkpointing активирован.")
+
+    criterion = nn.CrossEntropyLoss()
+
+    # Сначала переносим модель на ROCm/HIP устройство через accelerator
+    model = accelerator.prepare(model)
+
+    # === НАСТРОЙКА ADAFACTOR ДЛЯ GPU ===
+    optimizer = Adafactor(
+        model.parameters(),
+        lr=2e-4,  # Пиковый LR
+        weight_decay=0.01,
+        relative_step=False,
+        scale_parameter=False,
+        warmup_init=False
+    )
+
+    # === РАСЧЕТ И ИНИЦИАЛИЗАЦИЯ ПЛАНИРОВЩИКА С WARMUP ===
+    scheduler = None
+    if USE_COSINE_SCHEDULER:
+        # Считаем общее количество реальных шагов обновления весов (optimizer steps)
+        # 2000 строк в шарде / batch_size 1 / grad_accumulation_steps 4 = 500 шагов на шард.
+        steps_per_shard = math.ceil(2000 / (batch_size * grad_accumulation_steps))
+        total_steps = len(pt_files) * steps_per_shard
+
+        # Задаем warmup (например, 5% от общего числа шагов обучения)
+        num_warmup_steps = int(0.05 * total_steps)
+
+        # Комбинированный планировщик: плавно поднимает LR до 2e-4, затем опускает по косинусу до 0
+        scheduler = get_cosine_schedule_with_warmup(
+            optimizer=optimizer,
+            num_warmup_steps=num_warmup_steps,
+            num_training_steps=total_steps
+        )
+
+        if accelerator.is_local_main_process:
+            print(f"-> Планировщик АКТИВИРОВАН.")
+            print(f"   Всего шагов обучения: {total_steps}")
+            print(f"   Шагов разогрева (warmup): {num_warmup_steps}")
+
+    # Подготавливаем оптимизатор и планировщик через accelerator
+    if scheduler is not None:
+        optimizer, scheduler = accelerator.prepare(optimizer, scheduler)
+    else:
+        optimizer = accelerator.prepare(optimizer)
+
+    model.train()
+    shard_counter = 0
+
+    if accelerator.is_local_main_process:
+        print("Запуск обучения...")
+
+    for shard_path in pt_files:
+        shard_name = os.path.basename(shard_path)
+        if accelerator.is_local_main_process:
+            print(f"\n[Шард {shard_counter + 1}/{len(pt_files)}] {shard_name}")
+
+        shard_dataset = SingleShardDataset(shard_path)
+
+        train_loader = DataLoader(
+            shard_dataset,
+            batch_size=batch_size,
+            shuffle=True,
+            num_workers=2,
+            pin_memory=True
+        )
+
+        train_loader = accelerator.prepare(train_loader)
+
+        progress_bar = tqdm(
+            train_loader,
+            desc=f"Обработка {shard_name}",
+            disable=not accelerator.is_local_main_process
+        )
+
+        epoch_loss = 0.0
+        for step, batch in enumerate(progress_bar):
+            input_ids = batch
+            inputs = input_ids[:, :-1]
+            targets = input_ids[:, 1:]
+
+            with accelerator.accumulate(model):
+                logits, _ = model(inputs)
+                loss = criterion(logits.reshape(-1, logits.size(-1)), targets.reshape(-1))
+
+                accelerator.backward(loss)
+
+                optimizer.step()
+
+                # Делаем шаг планировщика только при реальном обновлении градиентов
+                if scheduler is not None and accelerator.sync_gradients:
+                    if not getattr(accelerator, "optimizer_step_was_skipped", False):
+                        scheduler.step()
+
+                optimizer.zero_grad()
+
+                epoch_loss += loss.item()
+                avg_loss = epoch_loss / (step + 1)
+                ppl = math.exp(avg_loss) if avg_loss < 20 else float('inf')
+
+                if accelerator.is_local_main_process:
+                    # Извлекаем текущий LR для вывода на панель tqdm
+                    current_lr = scheduler.get_last_lr()[0] if scheduler is not None else optimizer.param_groups[0]['lr']
+                    progress_bar.set_postfix({
+                        "loss": f"{loss.item():.4f}",
+                        "avg_loss": f"{avg_loss:.4f}",
+                        "ppl": f"{ppl:.1f}",
+                        "lr": f"{current_lr:.2e}"
+                    })
+
+        shard_counter += 1
+
+        # Сохранение после каждого шарда
+        if accelerator.is_local_main_process:
+            current_checkpoint_path = os.path.join(checkpoint_dir, f"jirack_shard_{shard_counter}")
+            os.makedirs(current_checkpoint_path, exist_ok=True)
+            unwrapped_model = accelerator.unwrap_model(model)
+            torch.save(unwrapped_model.state_dict(), os.path.join(current_checkpoint_path, "model_weights.pt"))
+            print(f"✅ Сохранено после шарда {shard_counter}")
+
+    # Финальное сохранение
+    accelerator.wait_for_everyone()
+    if accelerator.is_local_main_process:
+        final_dir = os.path.join(checkpoint_dir, "jirack_final_3b")
+        os.makedirs(final_dir, exist_ok=True)
+        unwrapped_model = accelerator.unwrap_model(model)
+        torch.save(unwrapped_model.state_dict(), os.path.join(final_dir, "model_final_weights.pt"))
+        print("Финальное сохранение завершено.")
+
+if __name__ == "__main__":
+    train()