feat: update training pipeline to use full-precision merge and add volume verification test

app.py

@@ -16,7 +16,7 @@ GRAMMAR_PATH = ROOT / "grammar.gbnf"
 
 # Keep the demo runnable while the GGUF files are still being prepared.
 MOCK_MODE = os.getenv("STEP_ZERO_MOCK", "1") != "0"
-NEMOTRON_MODEL_PATH = os.getenv("NEMOTRON_MODEL_PATH", "./models/nemotron-mini-4b.gguf")
+NEMOTRON_MODEL_PATH = os.getenv("NEMOTRON_MODEL_PATH", "./models/step-zero-nemotron-finetuned.gguf")
 MINICPM_MODEL_PATH = os.getenv("MINICPM_MODEL_PATH", "./models/minicpm-3-4b.gguf")
 
 app = FastAPI(title="Step-Zero")

modal_train.py

@@ -38,7 +38,7 @@ def train_model():
         model_name = "nvidia/Nemotron-Mini-4B-Instruct",
         max_seq_length = max_seq_length,
         dtype = None,
-        load_in_4bit = True,
+        load_in_4bit = False,
     )
 
     model = FastLanguageModel.get_peft_model(
@@ -106,21 +106,42 @@ def train_model():
     )
     trainer.train()
     
-    print("Exporting HuggingFace 16-bit Model...")
-    model.save_pretrained_merged("step-zero-nemotron-hf", tokenizer, save_method="merged_16bit")
+    print("Merging LoRA adapter into base model...")
+    import os, subprocess
+    hf_dir = "/root/step-zero-nemotron-hf"
+    os.makedirs(hf_dir, exist_ok=True)
     
-    import subprocess
-    print("Cloning llama.cpp to run conversion manually...")
-    subprocess.run(["git", "clone", "https://github.com/ggerganov/llama.cpp.git"])
-    subprocess.run(["pip", "install", "-r", "llama.cpp/requirements.txt"])
+    # Merge the LoRA weights into the base model using PEFT directly
+    merged_model = model.merge_and_unload()
+    merged_model.save_pretrained(hf_dir)
+    tokenizer.save_pretrained(hf_dir)
     
-    print("Converting to GGUF using llama.cpp...")
-    subprocess.run([
-        "python3", "llama.cpp/convert_hf_to_gguf.py", 
-        "step-zero-nemotron-hf", 
+    print(f"HF export contents: {os.listdir(hf_dir)}")
+    
+    print("Cloning llama.cpp for manual GGUF conversion...")
+    subprocess.run(["git", "clone", "--depth", "1", "https://github.com/ggerganov/llama.cpp.git", "/root/llama.cpp"], check=True)
+    subprocess.run(["pip", "install", "-r", "/root/llama.cpp/requirements.txt"], check=True)
+    
+    print("Converting HF model to GGUF...")
+    result = subprocess.run([
+        "python3", "/root/llama.cpp/convert_hf_to_gguf.py", 
+        hf_dir, 
         "--outfile", "/vol/step-zero-nemotron-finetuned.gguf", 
         "--outtype", "q8_0"
-    ])
+    ], capture_output=True, text=True)
+    print("STDOUT:", result.stdout[-2000:] if result.stdout else "")
+    print("STDERR:", result.stderr[-2000:] if result.stderr else "")
+    if result.returncode != 0:
+        print(f"convert_hf_to_gguf.py failed with code {result.returncode}")
+        print("Trying f16 fallback...")
+        result2 = subprocess.run([
+            "python3", "/root/llama.cpp/convert_hf_to_gguf.py", 
+            hf_dir, 
+            "--outfile", "/vol/step-zero-nemotron-finetuned.gguf", 
+            "--outtype", "f16"
+        ], capture_output=True, text=True)
+        print("STDOUT:", result2.stdout[-2000:] if result2.stdout else "")
+        print("STDERR:", result2.stderr[-2000:] if result2.stderr else "")
     
     try:
         vol.commit()

test_modal_ls.py

@@ -0,0 +1,16 @@
+import modal
+
+image = modal.Image.from_registry("ubuntu:22.04")
+
+app = modal.App("step-zero-test-ls")
+vol = modal.Volume.from_name("step-zero-volume")
+
+@app.function(image=image, volumes={"/vol": vol})
+def check_vol():
+    import os
+    print("Files in /vol:")
+    os.system("ls -la /vol")
+
+@app.local_entrypoint()
+def main():
+    check_vol.remote()