p-doom/crowd-code-dataset-0.2 huggingface dataset:
[terminal_command] TERMINAL ()
python -c "import json; d = json.load(open('/p/home/jusers/mahajan1/juwels/projects/envcomp/miles/tab-model-eval/data/eval/output/13034357/qwen8B_no_lora_bs_32_13034357_1000_generations.jsonl')); print(f'Keys: {d.keys()}'); print(f'First task results keys: {d[\"generation_results\"][0].keys()}'); [print(f'Task: {r[\"task_id\"]}, Response head: {repr(r[\"response_text\"][:50])}') for r in d[\"generation_results\"]] "
[terminal_command] TERMINAL ()
python -c "import json; d = json.load(open('/p/home/jusers/mahajan1/juwels/projects/envcomp/miles/tab-model-eval/data/eval/output/13034357/qwen8B_no_lora_bs_32_13034357_1000_generations.jsonl')); [print(f'Task: {r[\"task_id\"]}, First Sample Head: {repr(r[\"samples\"][0][\"response_text\"][:50])}') for r in d[\"generation_results\"]] "
[terminal_command] TERMINAL ()
source .venv/bin/activate && python -c "import json; from transformers import AutoTokenizer; tokenizer = AutoTokenizer.from_pretrained('/p/scratch/envcomp/mihir/huggingface/Qwen3-8B', trust_remote_code=True); d = json.load(open('/p/home/jusers/mahajan1/juwels/projects/envcomp/miles/tab-model-eval/data/eval/output/13034357/qwen8B_no_lora_bs_32_13034357_1000_generations.jsonl')); task = [r for r in d['generation_results'] if r['task_id'] == 'eval_terminal_invocation/0'][0]; print(f'Length: {len(tokenizer.apply_chat_template(task[\"context\"], tokenize=True))}')"
[tab] peek_checkpoint.py (python)
import torch
import torch.distributed.checkpoint as dist_cp
import os
import pickle

class UnpicklerWrapper(pickle.Unpickler):
    def find_class(self, mod_name, name):
        class DummyClass:
            def __init__(self, *args, **kwargs):
                pass
        if mod_name.startswith("megatron") or mod_name.startswith("glm") or mod_name.startswith("miles"):
            return DummyClass
        return super().find_class(mod_name, name)

class WrappedStorageReader(dist_cp.FileSystemReader):
    def read_metadata(self):
        path = self.fs.concat_path(self.path, ".metadata")
        with self.fs.create_stream(path, "rb") as metadata_file:
            metadata = UnpicklerWrapper(metadata_file).load()
        if getattr(metadata, "storage_meta", None) is None:
            metadata.storage_meta = dist_cp.StorageMeta()
        metadata.storage_meta.load_id = self.load_id
        if metadata.planner_data is None:
            metadata.planner_data = {}
        return metadata

class EmptyStateDictLoadPlanner(dist_cp.default_planner.DefaultLoadPlanner):
    def set_up_planner(self, state_dict, metadata=None, is_coordinator=False):
        for k, v in metadata.state_dict_metadata.items():
            if isinstance(v, dist_cp.metadata.TensorStorageMetadata):
                v = torch.empty(v.size, dtype=v.properties.dtype)
            state_dict[k] = v
        super().set_up_planner(state_dict, metadata, is_coordinator)

input_dir = "/p/scratch/envcomp/mihir/huggingface/shared_data/13034357/checkpoints/iter_0001000/model"
state_dict = {}
dist_cp.state_dict_loader._load_state_dict(
    state_dict,
    storage_reader=WrappedStorageReader(input_dir),
    planner=EmptyStateDictLoadPlanner(),
    no_dist=True,
)

print("Keys in state_dict:")
for k in list(state_dict.keys())[:20]:
    print(k)
[terminal_command] TERMINAL ()
source .venv/bin/activate && python -c 'import torch; import torch.distributed.checkpoint as dist_cp; import os; import pickle; class UnpicklerWrapper(pickle.Unpickler): def find_class(self, mod_name, name): class DummyClass: def __init__(self, *args, **kwargs): pass; return DummyClass; class WrappedStorageReader(dist_cp.FileSystemReader): def read_metadata(self): path = self.fs.concat_path(self.path, ".metadata"); with self.fs.create_stream(path, "rb") as metadata_file: metadata = UnpicklerWrapper(metadata_file).load(); return metadata; input_dir = "/p/scratch/envcomp/mihir/huggingface/shared_data/13034357/checkpoints/iter_0001000/model"; reader = WrappedStorageReader(input_dir); metadata = reader.read_metadata(); [print(k) for k in metadata.state_dict_metadata.keys() if "lm_head" in k or "norm" in k]'
[tab] check_merged_inference.py (python)
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_path = "/p/scratch/envcomp/mihir/huggingface/shared_data/13034357/checkpoints/iter_0001000_merged"
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.bfloat16, device_map="cpu", trust_remote_code=True)

prompt = "How do I list files in a directory?"
messages = [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user", "content": prompt}
]
text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(text, return_tensors="pt")

print("Generating...")
outputs = model.generate(**inputs, max_new_tokens=50)
print(tokenizer.decode(outputs[0]))
[tab] peek_filtered_keys.py (python)
import torch
import torch.distributed.checkpoint as dist_cp
import os
import pickle

class UnpicklerWrapper(pickle.Unpickler):
    def find_class(self, mod_name, name):
        class DummyClass:
            def __init__(self, *args, **kwargs):
                pass
        return DummyClass

class WrappedStorageReader(dist_cp.FileSystemReader):
    def read_metadata(self):
        path = self.fs.concat_path(self.path, ".metadata")
        with self.fs.create_stream(path, "rb") as metadata_file:
            metadata = UnpicklerWrapper(metadata_file).load()
        return metadata

input_dir = "/p/scratch/envcomp/mihir/huggingface/shared_data/13034357/checkpoints/iter_0001000/model"
reader = WrappedStorageReader(input_dir)
metadata = reader.read_metadata()

print("Filtered keys:")
for k in metadata.state_dict_metadata.keys():
    if "lm_head" in k or ".norm." in k or "model.norm." in k:
        print(k)
[tab] peek_checkpoint_metadata.py (python)
import torch
import torch.distributed.checkpoint as dist_cp
import os
import pickle

class UnpicklerWrapper(pickle.Unpickler):
    def find_class(self, mod_name, name):
        class DummyClass:
            def __init__(self, *args, **kwargs):
                pass
        if mod_name.startswith("megatron") or mod_name.startswith("glm") or mod_name.startswith("miles"):
            return DummyClass
        return super().find_class(mod_name, name)

class WrappedStorageReader(dist_cp.FileSystemReader):
    def read_metadata(self):
        path = self.fs.concat_path(self.path, ".metadata")
        with self.fs.create_stream(path, "rb") as metadata_file:
            metadata = UnpicklerWrapper(metadata_file).load()
        if getattr(metadata, "storage_meta", None) is None:
            metadata.storage_meta = dist_cp.StorageMeta()
        metadata.storage_meta.load_id = self.load_id
        if metadata.planner_data is None:
            metadata.planner_data = {}
        return metadata

input_dir = "/p/scratch/envcomp/mihir/huggingface/shared_data/13034357/checkpoints/iter_0001000/model"
reader = WrappedStorageReader(input_dir)
metadata = reader.read_metadata()

print("Keys in state_dict_metadata:")
for k in list(metadata.state_dict_metadata.keys())[:50]:
    print(k)
[content] slurm/jobs/mihir/juelich/crowd-pilot/eval/eval_miles_run_batch.sh (shellscript)
python tools/merge_fsdp_checkpoint.py \
    --checkpoint-dir /p/scratch/envcomp/mihir/huggingface/shared_data/13034357/checkpoints/iter_0002000 \
    --output-dir /p/scratch/envcomp/mihir/huggingface/shared_data/13034357/checkpoints/iter_0002000_merged \
    --hf-source-dir /p/scratch/envcomp/mihir/huggingface/Qwen3-8B
[tab] slurm/jobs/mihir/juelich/crowd-pilot/eval/eval_anta_claus.sh (shellscript)
# LAUNCH ANTA CLAUS EVAL JOBS

# 1.7B
sbatch /p/home/jusers/mahajan1/juwels/projects/envcomp/miles/slurm/jobs/mihir/juelich/crowd-pilot/eval/eval_miles_run_batch.sh qwen1.7B_no_lora_bs_8_13032805
sbatch /p/home/jusers/mahajan1/juwels/projects/envcomp/miles/slurm/jobs/mihir/juelich/crowd-pilot/eval/eval_miles_run_batch.sh qwen1.7B_no_lora_bs_16_13032803
# sbatch /p/home/jusers/mahajan1/juwels/projects/envcomp/miles/slurm/jobs/mihir/juelich/crowd-pilot/eval/eval_miles_run_batch.sh qwen1.7B_lora_bs_8_13032802
sbatch /p/home/jusers/mahajan1/juwels/projects/envcomp/miles/slurm/jobs/mihir/juelich/crowd-pilot/eval/eval_miles_run_batch.sh qwen1.7B_lora_bs_16_13032800
sbatch /p/home/jusers/mahajan1/juwels/projects/envcomp/miles/slurm/jobs/mihir/juelich/crowd-pilot/eval/eval_miles_run_batch.sh qwen1.7B_no_lora_bs_32_13032804

# 4B
sbatch /p/home/jusers/mahajan1/juwels/projects/envcomp/miles/slurm/jobs/mihir/juelich/crowd-pilot/eval/eval_miles_run_batch.sh qwen4B_no_lora_bs_8_13032814
sbatch /p/home/jusers/mahajan1/juwels/projects/envcomp/miles/slurm/jobs/mihir/juelich/crowd-pilot/eval/eval_miles_run_batch.sh qwen4B_lora_bs_8_13032811
sbatch /p/home/jusers/mahajan1/juwels/projects/envcomp/miles/slurm/jobs/mihir/juelich/crowd-pilot/eval/eval_miles_run_batch.sh qwen4B_lora_bs_16_13032809
sbatch /p/home/jusers/mahajan1/juwels/projects/envcomp/miles/slurm/jobs/mihir/juelich/crowd-pilot/eval/eval_miles_run_batch.sh qwen4B_no_lora_bs_16_13032812
sbatch /p/home/jusers/mahajan1/juwels/projects/envcomp/miles/slurm/jobs/mihir/juelich/crowd-pilot/eval/eval_miles_run_batch.sh qwen4B_no_lora_bs_32_13032813
sbatch /p/home/jusers/mahajan1/juwels/projects/envcomp/miles/slurm/jobs/mihir/juelich/crowd-pilot/eval/eval_miles_run_batch.sh qwen4B_lora_bs_32_13032810

# 8B
sbatch /p/home/jusers/mahajan1/juwels/projects/envcomp/miles/slurm/jobs/mihir/juelich/crowd-pilot/eval/eval_miles_run_batch.sh qwen8B_no_lora_bs_32_13034357
sbatch /p/home/jusers/mahajan1/juwels/projects/envcomp/miles/slurm/jobs/mihir/juelich/crowd-pilot/eval/eval_miles_run_batch.sh qwen8B_lora_bs_16_13032815
sbatch /p/home/jusers/mahajan1/juwels/projects/envcomp/miles/slurm/jobs/mihir/juelich/crowd-pilot/eval/eval_miles_run_batch.sh qwen8B_lora_bs_8_13032817
sbatch /p/home/jusers/mahajan1/juwels/projects/envcomp/miles/slurm/jobs/mihir/juelich/crowd-pilot/eval/eval_miles_run_batch.sh qwen8B_no_lora_bs_16_13032818
sbatch /p/home/jusers/mahajan1/juwels/projects/envcomp/miles/slurm/jobs/mihir/juelich/crowd-pilot/eval/eval_miles_run_batch.sh qwen8B_lora_bs_32_13032816
sbatch /p/home/jusers/mahajan1/juwels/projects/envcomp/miles/slurm/jobs/mihir/juelich/crowd-pilot/eval/eval_miles_run_batch.sh qwen8B_no_lora_bs_8_13032821

# 32B
sbatch /p/home/jusers/mahajan1/juwels/projects/envcomp/miles/slurm/jobs/mihir/juelich/crowd-pilot/eval/eval_miles_run_batch.sh qwen32B_no_lora_bs_32_13032806
sbatch /p/home/jusers/mahajan1/juwels/projects/envcomp/miles/slurm/jobs/mihir/juelich/crowd-pilot/eval/eval_miles_run_batch.sh qwen32B_lora_bs_8_13032807
[tab] slurm/jobs/mihir/juelich/crowd-pilot/2-wave-broken-resolutions/sbatch_template.sh (shellscript)
#!/usr/bin/env bash
#
# SBATCH Template Script for SFT Training
# This script is invoked by the launcher with environment variables set
#
# Required environment variables (set by launcher):
#   - MODEL_SIZE: 1.7b, 4b, or 8b
#   - BATCH_SIZE: 8, 16, or 32
#   - USE_LORA: true or false
#

# #SBATCH --nodes=4  # Set dynamically by launcher
#SBATCH --ntasks-per-node=1
#SBATCH --time=24:00:00
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:4
#SBATCH --mem=200GB
#SBATCH --account=envcomp
#SBATCH --nodelist=jwb[0700-0799]
#SBATCH --signal=USR1@60
#SBATCH --output=/p/project1/envcomp/mihir/logs/train/crowd-pilot-wave2-broken-resolutions/%x_%j.log
#SBATCH --error=/p/project1/envcomp/mihir/logs/train/crowd-pilot-wave2-broken-resolutions/%x_%j.log

trap 'echo "Job killed by signal"' USR1

# Print this script for debugging
cat $0

echo "============================================"
echo "SBATCH Environment:"
echo "  MODEL_SIZE: ${MODEL_SIZE}"
echo "  BATCH_SIZE: ${BATCH_SIZE}"
echo "  USE_LORA:   ${USE_LORA}"
echo "  CP_SIZE:    ${CONTEXT_PARALLEL_SIZE}"
echo "  SLURM_JOB_ID: ${SLURM_JOB_ID}"
echo "============================================"

# Load required modules
module load CUDA/12
module load cuDNN/9.5.0.50-CUDA-12
module load NCCL/default-CUDA-12
module load Clang/18.1.8

# Activate virtual environment
source .venv/bin/activate

# Export configuration for the run script
export MODEL_SIZE="${MODEL_SIZE}"
export BATCH_SIZE="${BATCH_SIZE}"
export USE_LORA="${USE_LORA}"
export CONTEXT_PARALLEL_SIZE="${CONTEXT_PARALLEL_SIZE}"

# Run the training script
srun sh slurm/jobs/mihir/juelich/crowd-pilot/2-wave-broken-resolutions/run_template.sh