p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] scripts/run-sft-torchrun.sh (shellscript)
#!/bin/bash
#
# Ray-free SFT Training Script
#
export PYTHONUNBUFFERED=1
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
# FIXME(f.srambical): this is hardcoded for now
GPUS_PER_NODE=${SLURM_GPUS_ON_NODE}
NUM_NODES=${SLURM_JOB_NUM_NODES}
NODE_RANK=${SLURM_NODEID}
MASTER_ADDR=${MASTER_ADDR:-$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n 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)"

export NCCL_DEBUG=INFO
export TORCH_DISTRIBUTED_DEBUG=INFO

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
   --save ${LOAD_SAVE_PATH}
   --save-interval 10
)

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

   --loss-type sft_loss
   --calculate-per-token-loss
   --disable-compute-advantages-and-returns
   --num-rollout 2000
)

LORA_ARGS=(
    # --use-lora
    # --lora-rank 8
    # --lora-alpha 16
    # --lora-dropout 0.0
    # --lora-target-modules q_proj v_proj
)

OPTIMIZER_ARGS=(
   --optimizer adam
   --lr 1e-5
   --lr-decay-style WSD
   --lr-wsd-decay-style linear
   --lr-warmup-iters 100
   --lr-decay-iters 2000
   --lr-wsd-decay-iters 500
   --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-torchrun
)

TRAIN_BACKEND_ARGS=(
   --train-backend fsdp
   --update-weight-buffer-size 536870912
   --gradient-checkpointing
   --attn-implementation flash_attention_3
)

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

MISC_ARGS=(
   --rollout-max-context-len 8192
   --rollout-max-prompt-len 8000
   --rollout-max-response-len 8192
   --dump-details /fast/project/HFMI_SynergyUnit/tab_model/huggingface/shared_data/qwen3-600M-fsdp-1116-noref/dump_details
)

torchrun \
    --nproc_per_node=${GPUS_PER_NODE} \
    --nnodes=${NUM_NODES} \
    --node_rank=${NODE_RANK} \
    --master_addr=${MASTER_ADDR} \
    --master_port=${MASTER_PORT:-29500} \
    train_sft.py \
    ${CKPT_ARGS[@]} \
    ${SFT_ARGS[@]} \
    ${LORA_ARGS[@]} \
    ${OPTIMIZER_ARGS[@]} \
    ${WANDB_ARGS[@]} \
    ${TRAIN_BACKEND_ARGS[@]} \
    ${PERF_ARGS[@]} \
    ${MISC_ARGS[@]}
[content] train_sft.py (python)
# Only optimize trainable parameters (important for LoRA to save memory)
        trainable_params = [p for p in self.model.parameters() if p.requires_grad]
        
        if self.args.use_lora:
            total_params = sum(p.numel() for p in self.model.parameters())
            trainable_count = sum(p.numel() for p in trainable_params)
            logger.info(
                f"[Rank {dist.get_rank()}] LoRA: {trainable_count:,} trainable params "
                f"out of {total_params:,} total ({100 * tra
[content] train_sft.py (python)
if self.args.use_lora:
                logger.warning(
                    "Gradient checkpointing is incompatible with LoRA (frozen base model "
                    "means no gradients flow through checkpointed inputs). Disabling gradient checkpointing."
                )
            else:
[content] train_sft.py (python)
# Re-apply requires_grad settings after FSDP wrapping
        # FSDP creates new parameter objects, so the settings from apply_lora are lost
        if self.args.use_lora:
            for n, p in self.model.named_parameters():
                if "lora_" in n:
                    p.requires_grad = True
                else:
                    p.requires_grad = False
[tab] /home/franz.srambical/slurm/dev/franz/berlin/crowd-pilot/nemo/qwen3_600m/test_lora_finetune_4k_bs_1.py (python)
import os

import nemo_run as run
from nemo.collections import llm

os.environ["NEMORUN_HOME"] = "/fast/project/HFMI_SynergyUnit/tab_model/franz/nemorun_home"
os.environ["NEMO_CACHE_DIR"] = "/fast/project/HFMI_SynergyUnit/tab_model/franz/nemo_cache"

NAME = "qwen3_600m_lora_finetune_crowd_code_dataset_4k_tokens_batch_size_1"
DATASET_DIR = "/fast/project/HFMI_SynergyUnit/tab_model/data/nemo_hf_part_jsonl_4k_tokens/" 

recipe = llm.qwen3_600m.finetune_recipe(
    name=NAME,
    dir=f"/fast/project/HFMI_SynergyUnit/tab_model/data/outputs",
    num_nodes=1,
    num_gpus_per_node=1,
    peft_scheme='lora',  # 'lora', 'none'
    packed_sequence=False,
)

def local_executor_torchrun(nodes: int, devices: int) -> run.LocalExecutor:
    # Env vars for jobs are configured here
    env_vars = {
        "TORCH_NCCL_AVOID_RECORD_STREAMS": "1",
        "NCCL_NVLS_ENABLE": "0",
        "NVTE_DP_AMAX_REDUCE_INTERVAL": "0",
        "NVTE_ASYNC_AMAX_REDUCTION": "1",
    }

    executor = run.LocalExecutor(ntasks_per_node=devices, launcher="torchrun", env_vars=env_vars)

    return executor

def run_finetuning():
    recipe.resume.restore_config.path = "/fast/project/HFMI_SynergyUnit/tab_model/data/checkpoints/nemo_converted_weights_qwen3-600m/"
    recipe.data = run.Config(
        llm.ChatDataModule,
        dataset_root=DATASET_DIR,
        seq_length=4096,
        global_batch_size=1,
        micro_batch_size=1,
    )

    recipe.trainer.strategy.tensor_model_parallel_size = 1
    recipe.trainer.strategy.expert_model_parallel_size = 1

    executor = local_executor_torchrun(nodes=recipe.trainer.num_nodes, devices=recipe.trainer.devices)

    run.run(recipe, executor=executor, name=NAME)

# This condition is necessary for the script to be compatible with Python's multiprocessing module.
if __name__ == "__main__":
    run_finetuning()