p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] slurm/jobs/mihir/berlin/crowd-pilot/run_sqwen3-0.6b-sft-torch-crowd-code.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=24:00:00
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:2
#SBATCH --mem=600GB
#SBATCH --signal=USR1@60
#SBATCH --output=/fast/project/HFMI_SynergyUnit/mihir/logs/train/qwen3-0.6b-sft-torch-crowd-code/%x_%j.log
#SBATCH --error=/fast/project/HFMI_SynergyUnit/mihir/logs/train/qwen3-0.6b-sft-torch-crowd-code/%x_%j.log
#SBATCH --job-name=run_qwen3-0.6b-sft-torch-crowd-code


trap 'echo "Job killed by signal"' USR1

cat $0
module load CUDA/12.8

source .venv/bin/activate

srun sh scripts/run-qwen3-0.6B-torch-sft-crowd-code.sh
[tab] scripts/run-qwen3-0.6B-torch-sft-pythonic.sh (shellscript)
#!/bin/bash

# for rerun the task
pkill -9 sglang
sleep 3
ray stop --force
pkill -9 ray
pkill -9 python
sleep 3
pkill -9 ray
pkill -9 python

# set -ex

# will prevent ray from buffering stdout/stderr
export PYTHONUNBUFFERED=1

NVLINK_COUNT=$(nvidia-smi | grep -o "NVLink" | wc -l)
if [ "$NVLINK_COUNT" -gt 0 ]; then
    HAS_NVLINK=1
else
    HAS_NVLINK=0
fi
echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)"


# --- 1. DYNAMIC HOST IP DETECTION (CRITICAL FOR SLURM) ---
# Don't hardcode IP. Get the actual IP of the current node.
# export HEAD_NODE_IP=$(hostname -I | awk '{print $1}')
# echo "Detected Head Node IP: ${HEAD_NODE_IP}"

# # --- 2. PROXY CONFIGURATION ---
# # Ensure local traffic doesn't go through a corporate proxy
# export no_proxy="${HEAD_NODE_IP},localhost,127.0.0.1,0.0.0.0"
# export NO_PROXY="${HEAD_NODE_IP},localhost,127.0.0.1,0.0.0.0"

SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)"

RUN_ID=${RUN_ID:-"run_$(date +%Y%m%d_%H%M%S)"}
LOAD_SAVE_PATH="/fast/project/HFMI_SynergyUnit/tab_model/huggingface/shared_data/${RUN_ID}/checkpoints"

CKPT_ARGS=(
   --hf-checkpoint /fast/project/HFMI_SynergyUnit/tab_model/huggingface/Qwen3-0.6B
   --load /fast/project/HFMI_SynergyUnit/tab_model/huggingface/Qwen3-0.6B
   --ref-load /fast/project/HFMI_SynergyUnit/tab_model/huggingface/Qwen3-0.6B
)

SFT_ARGS=(
   --rollout-function-path miles.rollout.sft_rollout.generate_rollout
   --prompt-data /fast/project/HFMI_SynergyUnit/tab_model/huggingface/openhermes2_5.parquet
   --input-key messages
   --apply-chat-template
   --rollout-shuffle
   --num-epoch 3
   --rollout-batch-size 128
   --global-batch-size 128

   --loss-type sft_loss
   --calculate-per-token-loss
   --disable-compute-advantages-and-returns
   --debug-train-only
)

OPTIMIZER_ARGS=(
   --optimizer adam
   --lr 1e-5
   --lr-decay-style constant
   --weight-decay 0.1
   --adam-beta1 0.9
   --adam-beta2 0.98
)

WANDB_ARGS=(
   --use-wandb
   --wandb-project crowd-pilot-miles
   --wandb-team instant-uv
   --wandb-group qwen3-0.6b-sft-torch
)

SGLANG_ARGS=(
   --rollout-num-gpus-per-engine 2
   --sglang-mem-fraction-static 0.5
   --sglang-decode-log-interval 1000
   --sglang-chunked-prefill-size 4096
   --sglang-attention-backend fa3
)

TRAIN_BACKEND_ARGS=(
   --train-backend fsdp
   --update-weight-buffer-size 536870912
   --gradient-checkpointing
   --attn-implementation flash_attention_3
   --train-env-vars '{"PYTORCH_CUDA_ALLOC_CONF":"expandable_segments:True"}'
   --actor-num-gpus-per-node 2
)

PERF_ARGS=(
   --use-dynamic-batch-size
   --max-tokens-per-gpu 9216
)

# Build runtime env
RUNTIME_ENV_JSON="{
  \"env_vars\": {
    \"PYTHONPATH\": \"/fast/project/HFMI_SynergyUnit/mihir/Megatron-LM/\",
    \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\",
    \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\",
    \"PYTORCH_CUDA_ALLOC_CONF\": \"expandable_segments:True\"
  }
}"

python3 train.py \
   ${CKPT_ARGS[@]} \
   ${SFT_ARGS[@]} \
   ${OPTIMIZER_ARGS[@]} \
   ${WANDB_ARGS[@]} \
   ${SGLANG_ARGS[@]} \
   ${TRAIN_BACKEND_ARGS[@]} \
   ${PERF_ARGS[@]} \
   ${MISC_ARGS[@]}
[tab] miles/ray/ray_actor.py (python)
from miles.utils.misc import get_current_node_ip, get_free_port


class RayActor:
    @staticmethod
    def _get_current_node_ip_and_free_port(start_port=10000, consecutive=1):
        return get_current_node_ip(), get_free_port(start_port=start_port, consecutive=consecutive)

    def get_master_addr_and_port(self):
        return self.master_addr, self.master_port
[tab] examples/multi_agent/prompts.py (python)
## 定义了一些 prompts 的模板，用于生成不同的 prompts


SOLVER_PROMPT_TEMPLATE = """{problem_statement}"""


def generate_rewriter_template(num_solutions: int) -> str:
    """Dynamically generate rewrite templates based on the number of solutions."""
    solution_sections = []
    for i in range(num_solutions):
        solution_sections.append(f"#### Solution {i+1}
{{solution{i+1}}}

---")

    solutions_text = "
".join(solution_sections)

    return f"""### Task: Solution Rewriting Based on Previous Solutions ###
You are being reactivated to revise your mathematical proof. You are provided with two documents:
1.  The problem you need to solve.
2.  Your {num_solutions} different "Previous Solutions".

Your sole task is to generate a new, correct version of your solution based on your previous discoveries in the provided {num_solutions} solutions.

Refer to the following {num_solutions} solutions and solve the problem.
---

### Problem

{{problem_statement}}

---

### Candidates Solution
{solutions_text}
"""


def generate_select_template(num_solutions: int) -> str:
    """Dynamically generate select templates based on the number of solutions."""
    solution_sections = []
    for i in range(num_solutions):
        solution_sections.append(f"#### Solution {i+1}
{{solution{i+1}}}

---")

    solutions_text = "
".join(solution_sections)

    return f"""You will be given a challenging math problem followed by {num_solutions} solutions.
Your task is to systematically analyze these solutions to identify the most mathematically sound approach. 

You are provided with two documents:
1.  The problem you need to solve.
2.  Your {num_solutions} "Candidate Solutions".

Evaluation Process:
1. Initial Screening
- Group solutions by their final answers
- Identify and explain mathematical contradictions between different answers
- Eliminate solutions with clear mathematical errors

2. Detailed Analysis
For remaining solutions, evaluate:
- Mathematical precision and accuracy
- Logical progression of steps
- Completeness of mathematical reasoning
- Handling of edge cases or special conditions
- For solutions containing and addressing errors, evaluate the error identification and correction methodology.

3. Solution Comparison
Compare viable solutions based on:
- Efficiency of approach
- Clarity of mathematical reasoning
- Sophistication of method
- Robustness of solution (works for all cases)

Your response should include:
1. Brief analysis of conflicting answers
2. Detailed evaluation of mathematically sound solutions
3. Justification for eliminating incorrect solutions
4. Clear explanation for selecting the best approach

End your evaluation with exactly:
Judgment: IDX
where IDX is the index 1-{num_solutions} of the best solution

### Problem

{{problem_statement}}

---

### Candidate Solutions
{solutions_text}
"""
[tab] miles/rollout/sft_rollout.py (python)
import logging

from transformers import AutoTokenizer

from miles.utils.mask_utils import MultiTurnLossMaskGenerator

__all__ = ["generate_rollout"]

logger = logging.getLogger(__name__)


TOKENIZER = None
MASK_GENERATOR = None
SAMPLE_PRINTED = False


def generate_rollout(args, rollout_id, data_buffer, evaluation=False):
    """An example to implement the generate_rollout function for an rule based rm rollout generation.

    Args:
        args: the whole args
        rollout_id: int, the id of the rollout, used for deterministic data generation
        data_buffer: the data buffer to store the generated samples
        evaluation: bool, whether the rollout is for evaluation or not

    Returns:
        list[Sample]: a list of samples generated by the rollout
    """
    assert not evaluation
    assert args.rollout_global_dataset

    global TOKENIZER, MASK_GENERATOR, SAMPLE_PRINTED
    if TOKENIZER is None:
        TOKENIZER = AutoTokenizer.from_pretrained(args.hf_checkpoint, trust_remote_code=True)

    if MASK_GENERATOR is None:
        MASK_GENERATOR = MultiTurnLossMaskGenerator(TOKENIZER, tokenizer_type=args.loss_mask_type)

    samples = data_buffer.get_samples(args.rollout_batch_size)

    for i, sample in enumerate(samples):
        (sample,) = sample
        messages = sample.prompt
        token_ids, loss_mask = MASK_GENERATOR.get_loss_mask(messages)
        response_length = MASK_GENERATOR.get_response_lengths([loss_mask])[0]

        sample.tokens = token_ids
        sample.response_length = response_length
        sample.reward = 0
        sample.loss_mask = loss_mask[-response_length:]

        if i == 0 and not SAMPLE_PRINTED:
            logger.info(
                f"sft_rollout::generate_rollout example data: {sample=} (raw){messages=} (raw){token_ids=} (raw){loss_mask=} {response_length=}"
            )
            SAMPLE_PRINTED = True

    return samples
[tab] /fast/home/mihir.mahajan/Projects/miles/miles/rollout/sft_rollout.py (python)
import logging

from transformers import AutoTokenizer

from miles.utils.mask_utils import MultiTurnLossMaskGenerator

__all__ = ["generate_rollout"]

logger = logging.getLogger(__name__)


TOKENIZER = None
MASK_GENERATOR = None
SAMPLE_PRINTED = False


def generate_rollout(args, rollout_id, data_buffer, evaluation=False):
    """An example to implement the generate_rollout function for an rule based rm rollout generation.

    Args:
        args: the whole args
        rollout_id: int, the id of the rollout, used for deterministic data generation
        data_buffer: the data buffer to store the generated samples
        evaluation: bool, whether the rollout is for evaluation or not

    Returns:
        list[Sample]: a list of samples generated by the rollout
    """
    assert not evaluation
    assert args.rollout_global_dataset

    global TOKENIZER, MASK_GENERATOR, SAMPLE_PRINTED
    if TOKENIZER is None:
        TOKENIZER = AutoTokenizer.from_pretrained(args.hf_checkpoint, trust_remote_code=True)

    if MASK_GENERATOR is None:
        MASK_GENERATOR = MultiTurnLossMaskGenerator(TOKENIZER, tokenizer_type=args.loss_mask_type)

    samples = data_buffer.get_samples(args.rollout_batch_size)

    for i, sample in enumerate(samples):
        (sample,) = sample
        messages = sample.prompt
        token_ids, loss_mask = MASK_GENERATOR.get_loss_mask(messages)
        response_length = MASK_GENERATOR.get_response_lengths([loss_mask])[0]

        sample.tokens = token_ids
        sample.response_length = response_length
        sample.reward = 0
        sample.loss_mask = loss_mask[-response_length:]

        if i == 0 and not SAMPLE_PRINTED:
            logger.info(
                f"sft_rollout::generate_rollout example data: {sample=} (raw){messages=} (raw){token_ids=} (raw){loss_mask=} {response_length=}"
            )
            SAMPLE_PRINTED = True

    return samples
[tab] slurm/jobs/mihir/berlin/crowd-pilot/convert_nemo_dataset.py (python)
from datasets import load_dataset
import os

# Define input and output paths
input_path = "/fast/project/HFMI_SynergyUnit/tab_model/data/nemo_hf_part_jsonl_4k_tokens/training.jsonl"
# Create output directory if it doesn't exist, though the path seems to indicate it should exist
output_path = "/fast/project/HFMI_SynergyUnit/tab_model/huggingface/nemo_hf_part_jsonl_4k_tokens.parquet"
os.makedirs(os.path.dirname(output_path), exist_ok=True)

print(f"Loading dataset from {input_path}")
ds = load_dataset("json", data_files=input_path)["train"]

def convert(sample):
    messages = []
    
    # Handle top-level system prompt
    if sample.get("system"):
        messages.append({
            "role": "system",
            "content": sample["system"]
        })

    conversations = sample["conversations"]

    def convert_role(role):
        # Map roles from the input format to the target format
        if role == "User":
            return "user"
        elif role == "Assistant":
            return "assistant"
        elif role == "System":
            return "system"
        # Fallback for lowercase or other variations if necessary, though based on inspection they are Title Case
        elif role.lower() == "user":
            return "user"
        elif role.lower() == "assistant":
            return "assistant"
        elif role.lower() == "system":
            return "system"
        else:
            raise ValueError(f"Unknown role: {role}")

    for turn in conversations:
        messages.append({
            "role": convert_role(turn["from"]),
            "content": turn["value"],
        })

    return {"messages": messages}

print("Converting dataset...")
ds = ds.map(convert, remove_columns=ds.column_names)
print(f"Dataset converted. Saving to {output_path}")
ds.to_parquet(output_path)
print("Done.")
[content] /fast/home/mihir.mahajan/Projects/miles/miles/ray/rollout.py (python)
                print("="*100)
                print("Evaluation data")
                print(type(data))
                print(data[0])
                print(data[0][0])
                print("="*100)
                dump_data = dict(
                    samples=[sample.to_dict() for dataset_name, info in data.items() for sample in info["samples"]]
[content] /fast/home/mihir.mahajan/Projects/miles/miles/ray/rollout.py (python)
                if isinstance(data, list):
                    # Data is List[List[Sample]] or List[Sample]
                    # We flatten the list structure safely
                    all_samples = []
                    for item in data:
                        if isinstance(item, list):
                            all_samples.extend(item)
                        else:
                            all_samples.append(item)
                            
                    dump_data = dict(
                        samples=[sample.to_dict() for sample in all_samples]
                    )
                else:
                    # Original logic for Dictionary structure
                    dump_data = dict(
                        samples=[sample.to_dict() for dataset_name, info in data.items() for sample in info["samples"]]
                    )
[content] /fast/home/mihir.mahajan/Projects/miles/miles/ray/rollout.py (python)
    if isinstance(data, list):
        # 1. Flatten the nested list (List[List[Sample]] -> List[Sample])
        all_samples = []
        for item in data:
            if isinstance(item, list):
                all_samples.extend(item)
            else:
                all_samples.append(item)
        
        # 2. Extract the metrics that the loop below expects (rewards, truncated)
        #    We use the same logic as in _convert_samples_to_train_data
        rewards = [s.get_reward_value(args) for s in all_samples]
        truncated = [1 if s.status == Sample.Status.TRUNCATED else 0 for s in all_samples]
        
        # 3. Wrap it in a dictionary with a default key
        data = {
            "default": {
                "samples": all_samples,
                "rewards": rewards,
                "truncated": truncated
            }
        }
[content] miles/backends/fsdp_utils/actor.py (python)
        """Run one training update over a rollout batch.

        Parameters:
            rollout_id: Monotonic id for logging.
            rollout_data_ref: A Box handle wrapping a Ray object reference to a
                dictionary with rollout tensors and metadata (e.g., `tokens`,
                `loss_masks`, `rewards`, `response_lengths`, optional
                `rollout_log_probs`, etc.). It will be fetched and partitioned
                by `process_rollout_data` based on data-parallel rank/size.
        """
        if self.args.offload_train:
            self.wake_up()

        with inverse_timer("train_wait"), timer("train"):
            rollout_data = process_rollout_data(self.args, rollout_data_ref, self.dp_rank, self.dp_size)
            if self.args.debug_rollout_only:
                return
            self._train_core(rollout_id=rollout_id, rollout_data=rollout_data)

        train_metric_utils.log_perf_data_raw(
            rollout_id=rollout_id,
            args=self.args,
            is_primary_rank=dist.get_rank() == 0,
            compute_total_fwd_flops=None,
[content] miles/backends/fsdp_utils/actor.py (python)
    def _train_core(self, rollout_id: int, rollout_data) -> None:
        if self.args.advantage_estimator in ["grpo", "gspo"]:
            rollout_data["advantages"] = rollout_data["returns"] = [
                torch.tensor([rollout_data["rewards"][i]] * rollout_data["response_lengths"][i])
                for i in range(len(rollout_data["rewards"]))
            ]
        else:
            raise NotImplementedError(f"Unsupported advantage_estimator {self.args.advantage_estimator}")

        packed_batches, grad_accum = self._packed_data(rollout_data)

        assert (
            len(grad_accum) > 0
        ), f"Invalid grad_accum {grad_accum} for micro_batch_size {self.args.micro_batch_size} and global_batch_size {self.args.global_batch_size}"

        if self.ref_model is not None:
            self._compute_log_prob("ref", packed_batches, store_prefix="ref_")

        self._compute_log_prob("actor", packed_batches)
        self._log_rollout_data(rollout_id, rollout_data, packed_batches)

        with timer("actor_train"):
            reported_accum: dict[str, list[torch.Tensor]] = {}
            self.optimizer.zero_grad(set_to_none=True)
            for mbs_id, packed_batch in self.prof.iterate_train_actor(
                enumerate(tqdm(packed_batches, desc="actor_train", disable=dist.get_rank() != 0))
            ):
                self._train_step(
                    packed_batch=packed_batch,
                    reported_accum=reported_accum,
                    mbs_id=mbs_id,
                    grad_accum=grad_accum,
                )

        self.prof.step(rollout_id=rollout_id)

        train_dump_utils.save_debug_train_data(self.args, rollout_id=rollout_id, rollout_data=rollout_data)

        # Update ref model if needed (copy actor weights to ref)
        if (
            self.args.ref_update_interval is not None
            and (rollout_id + 1) % self.args.ref_update_interval == 0
            and self.ref_model is not None
        ):
            if dist.get_rank() == 0:
                logger.info(f"Updating ref model at rollout_id {rollout_id}")
            # Copy actor model state to ref model
            actor_state = self.model.state_dict()
            self.ref_model.load_state_dict(actor_state)
            self.ref_model.cpu()
[content] miles/backends/fsdp_utils/actor.py (python)

    def _train_core(self, rollout_id: int, rollout_data) -> None:
        if self.args.advantage_estimator in ["grpo", "gspo"]:
            rollout_data["advantages"] = rollout_data["returns"] = [
                torch.tensor([rollout_data["rewards"][i]] * rollout_data["response_lengths"][i])
                for i in range(len(rollout_data["rewards"]))
            ]
        else:
            raise NotImplementedError(f"Unsupported advantage_estimator {self.args.advantage_estimator}")

        packed_batches, grad_accum = self._packed_data(rollout_data)

        assert (
            len(grad_accum) > 0
        ), f"Invalid grad_accum {grad_accum} for micro_batch_size {self.args.micro_batch_size} and global_batch_size {self.args.global_batch_size}"

        if self.ref_model is not None:
            self._compute_log_prob("ref", packed_batches, store_prefix="ref_")

        self._compute_log_prob("actor", packed_batches)
        self._log_rollout_data(rollout_id, rollout_data, packed_batches)

        with timer("actor_train"):
            reported_accum: dict[str, list[torch.Tensor]] = {}
            self.optimizer.zero_grad(set_to_none=True)
            for mbs_id, packed_batch in self.prof.iterate_train_actor(
                enumerate(tqdm(packed_batches, desc="actor_train", disable=dist.get_rank() != 0))
            ):
                self._train_step(
                    packed_batch=packed_batch,
                    reported_accum=reported_accum,
                    mbs_id=mbs_id,
                    grad_accum=grad_accum,
                )

        self.prof.step(rollout_id=rollout_id)

        train_dump_utils.save_debug_train_data(self.args, rollout_id=rollout_id, rollout_data=rollout_data)

        # Update ref model if needed (copy actor weights to ref)
        if (
            self.args.ref_update_interval is not None
            and (rollout_id + 1) % self.args.ref_update_interval == 0
            and self.ref_model is not None
        ):
            if dist.get_rank() == 0:
                logger.info(f"Updating ref model at rollout_id {rollout_id}")
            # Copy actor model state to ref model
            actor_state = self.model.state_dict()
            self.ref_model.load_state_dict(actor_state)
            self.ref_model.cpu()
[content] train.py (python)

            if args.loss_type == "sft":
                ray.get(actor_model.async_val.remote(rollout_id, rollout_data_ref))
            else:
                ray.get(rollout_manager.eval.remote(rollout_id))
[content] train.py (python)
            if args.loss_type == "sft":
                ray.get(actor_model.async_val.remote(rollout_id, rollout_data_ref))
            else:
                ray.get(rollout_manager.eval.remote(rollout_id))