Create chat_405b_deepspeed.py

chat_405b_deepspeed.py

@@ -0,0 +1,90 @@
+#%%writefile chat_405b.py
+# ==============================================================================
+# COPYRIGHT (C) 2025-2026 KONSTANTIN VLADIMIROVICH GRABKO. ALL RIGHTS RESERVED.
+# JiRack 405B Chat with DeepSpeed ZeRO-3 Offload
+# ==============================================================================
+
+import torch
+import os
+import deepspeed
+from transformers import AutoTokenizer
+from JiRackTernaryPyTorch_405b import JiRackTernaryModel
+
+# --- ПУТИ ---
+LOCAL_PATH = "/content/drive/MyDrive/JiRack_405B_SFT_Steps_2/step_4500"
+TOKENIZER_ID = "meta-llama/Meta-Llama-3.1-405B-Instruct"
+
+# --- DEEPSPEED CONFIG ---
+ds_config = {
+    "fp16": {"enabled": True},
+    "zero_optimization": {
+        "stage": 3,
+        "offload_param": {"device": "cpu", "pin_memory": True},
+        "overlap_comm": True,
+        "contiguous_gradients": True,
+        "stage3_max_live_parameters": 1e8,
+        "stage3_max_reuse_distance": 1e8,
+        "stage3_prefetch_bucket_size": 5e7,
+        "stage3_param_persistence_threshold": 1e5,
+    },
+    "steps_per_print": 2000,
+    "train_batch_size": 1,
+}
+
+def main():
+    print("[*] Loading tokenizer...")
+    tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_ID)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    print("[*] Initializing JiRack 405B with DeepSpeed ZeRO-3...")
+    
+    # Инициализация модели внутри контекста DeepSpeed Zero
+    # Это позволяет загружать веса по частям, не забивая RAM мгновенно
+    with deepspeed.zero.Init(config_dict_or_path=ds_config):
+        model = JiRackTernaryModel.from_pretrained(
+            LOCAL_PATH,
+            torch_dtype=torch.float16,
+            trust_remote_code=True
+        )
+
+    # Инициализация движка DeepSpeed
+    ds_engine = deepspeed.initialize(model=model, config_params=ds_config)[0]
+    ds_engine.eval()
+
+    print("\n✅ JiRack 405B is ready (ZeRO-3 Offload Active).")
+    print("Type 'exit' to quit.\n")
+
+    while True:
+        try:
+            user_input = input("User: ")
+            if user_input.lower() in ["exit", "quit"]: break
+            if not user_input.strip(): continue
+
+            print("JiRack thinking...", end="\r")
+            
+            # Форматируем промпт
+            prompt = f"<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n\n{user_input}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+            inputs = tokenizer(prompt, return_tensors="pt").to(torch.cuda.current_device())
+
+            with torch.no_grad():
+                # Используем .module для доступа к оригинальному методу generate
+                outputs = ds_engine.module.generate(
+                    **inputs,
+                    max_new_tokens=128,
+                    do_sample=True,
+                    temperature=0.7,
+                    top_p=0.9,
+                    eos_token_id=[tokenizer.eos_token_id, 128001, 128008, 128009]
+                )
+
+            response = tokenizer.decode(outputs[0][inputs['input_ids'].shape[-1]:], skip_special_tokens=True)
+            print(f"JiRack: {response}\n")
+
+        except KeyboardInterrupt:
+            break
+        except Exception as e:
+            print(f"\n[!] Error: {e}")
+
+if __name__ == "__main__":
+    main()