Create sft_tain_no_offload.py

prepared_sft_data/sft_tain_no_offload.py

@@ -0,0 +1,121 @@
+# ==============================================================================
+# COPYRIGHT (C) 2025-2026 KONSTANTIN VLADIMIROVICH GRABKO. ALL RIGHTS RESERVED.
+# PATENT PENDING | CMS MANHATTAN JIRACK TECHNOLOGY
+# ==============================================================================
+
+import os, sys, torch, glob, gc, json
+from torch.utils.data import DataLoader
+from safetensors.torch import load_file, save_file
+from tqdm import tqdm
+from torch.cuda.amp import autocast
+
+sys.path.append("/content")
+from JiRackTernaryPyTorch_70b import JiRackTernaryConfig, JiRackTernary70B, JiRackFullLinear
+
+# --- CONFIG ---
+LR = 3e-6
+GRAD_ACCUM_STEPS = 128
+SAVE_DIR = "/content/full_checkpoints_70b"
+os.makedirs(SAVE_DIR, exist_ok=True)
+
+def save_full_trainable_parameters(model, path):
+    state_dict = {n: p.data.cpu() for n, p in model.named_parameters() if p.requires_grad}
+    save_file(state_dict, path)
+    print(f"\n[Checkpoint] Сохранено: {path}")
+
+def unpack_2bit(packed, shape, scale):
+    p = packed.to(torch.int32)
+    t = torch.stack([(p >> 6) & 3, (p >> 4) & 3, (p >> 2) & 3, p & 3], -1).view(-1)
+    return (t.to(torch.bfloat16) - 1.0)[:shape.numel()].view(tuple(shape.tolist())) * scale
+
+def setup_70b_pytorch():
+    print("🛠️ Сборка модели JiRack 70B (Pure PyTorch Mode)...")
+    with open("/content/JiRack_BitNet_70B_Packed/config.json", "r") as f:
+        config_dict = json.load(f)
+        config = JiRackTernaryConfig.from_dict(config_dict)
+    
+    with torch.device("meta"): 
+        model = JiRackTernary70B(config)
+    
+    shards = sorted(glob.glob("/content/JiRack_BitNet_70B_Packed/checkpoints/checkpoint-260000/*.safetensors"))
+    for shard in tqdm(shards, desc="Loading Shards"):
+        sd = load_file(shard, device="cpu")
+        for name, param in sd.items():
+            target = name.replace("model.", "")
+            parts = target.split('.')
+            if "weight" in name and any(x in name for x in ["proj", "ffn", "layers"]):
+                scale_key = name.replace("weight", "scale")
+                if scale_key in sd:
+                    shape = sd[name.replace("weight", "shape")]
+                    w = unpack_2bit(sd[name], shape, sd[scale_key]).to("cuda")
+                    curr = model
+                    for p in parts[:-2]: curr = getattr(curr, p)
+                    setattr(curr, parts[-2], JiRackFullLinear(w, shape, sd[scale_key].to("cuda")))
+                    continue
+            try:
+                curr = model
+                for p in parts[:-1]: curr = getattr(curr, p)
+                setattr(curr, parts[-1], torch.nn.Parameter(param.to(device="cuda", dtype=torch.bfloat16)))
+            except: continue
+        del sd; gc.collect()
+    
+    print("⏳ Материализация оставшихся компонентов на  CUDA...")
+    model.to_empty(device="cuda")
+    
+    # Решаем проблему с Gradient Checkpointing
+    print("🚀 Попытка активации  Gradient Checkpointing...")
+    try:
+        # Пытаемся активировать стандартным методом
+        model.gradient_checkpointing_enable()
+    except ValueError:
+        # Если Transformers ругается, форсируем атрибут напрямую, 
+        # если логика внутри модели завязана на флаг 'gradient_checkpointing'
+        if hasattr(model, "config"):
+            model.config.gradient_checkpointing = True
+            print("⚠️ Стандартный метод не подошел, активирован флаг в config.")
+        else:
+            print("❌ Модель не поддерживает Gradient Checkpointing. Обучение 70B может упасть с  OOM.")
+
+    return model
+
+if __name__ == "__main__":
+    model = setup_70b_pytorch()
+    
+    optimizer = torch.optim.AdamW(model.parameters(), lr=LR, betas=(0.9, 0.95), weight_decay=0.1)
+    criterion = torch.nn.CrossEntropyLoss(ignore_ind
+
+    shards = sorted(glob.glob("/content/JiRack_BitNet_70B_Packed/prepared_sft_data/*.pt"))
+    global_step = 0
+
+    for shard in shards:
+        if not os.path.exists(shard): continue
+        dataset = torch.load(shard, map_location='cpu', weights_only=False)
+        dataloader = DataLoader(dataset, batch_size=1, shuffle=True)
+        pbar = tqdm(dataloader, desc=f"Training {os.path.basename(shard)}")
+        
+        model.train()
+        for step, batch in enumerate(pbar):
+            ids = batch['input_ids'].to("cuda")
+            mask = batch['attention_mask'].to("cuda")
+            labels = ids.clone(); labels[mask == 0] = -100
+            
+            with autocast(dtype=torch.bfloat16):
+                out = model(ids)
+                logits = out.logits[:, :-1, :].contiguous().view(-1, config_dict.get("vocab_size", 128256))
+                loss = criterion(logits, labels[:, 1:].contiguous().view(-1))
+                loss = loss / GRAD_ACCUM_STEPS
+
+            loss.backward()
+
+            if (step + 1) % GRAD_ACCUM_STEPS == 0:
+                optimizer.step()
+                optimizer.zero_grad()
+                global_step += 1
+                pbar.set_postfix({"loss": f"{loss.item() * GRAD_ACCUM_STEPS:.4f}", "step": global_step})
+
+                if global_step % 200 == 0:
+                    save_path = os.path.join(SAVE_DIR, f"jirack_70b_pt_step_{global_step}.safetensors")
+                    save_full_trainable_parameters(model, save_path)
+            
+            if step % 20 == 0:
+                torch.cuda.empty_cache()