p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] miles/backends/fsdp_utils/checkpoint.py (python)
from torch.distributed.checkpoint.state_dict import (
    StateDictOptions,
    get_model_state_dict,
    get_state_dict,
    set_state_dict,
)
from torch.distributed.checkpoint.stateful import Stateful
[content] train_sft.py (python)
            if self.args.use_lora:
                lora_config = LoRAConfig(
                    lora_rank=self.args.lora_rank,
                    lora_alpha=self.args.lora_alpha,
                    lora_dropout=self.args.lora_dropout,
                    target_modules=self.args.lora_target_modules,
                )
                model = apply_lora(model, lora_config)
                logger.info(f"[Rank {dist.get_rank()}] Applied LoRA: {lora_config}")
[content] train_sft.py (python)
            # Use non-reentrant mode for gradient checkpointing (required for PEFT/LoRA)
            gc_kwargs = {"use_reentrant": False} if is_lora_model(self.model) else {}
            self.model.gradient_checkpointing_enable(gradient_checkpointing_kwargs=gc_kwargs)
[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_PATH="/fast/project/HFMI_SynergyUnit/tab_model/huggingface/Qwen3-0.6B"
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 ${LOAD_PATH}
   --ref-load /fast/project/HFMI_SynergyUnit/tab_model/huggingface/Qwen3-0.6B
   --save ${SAVE_PATH}
   --save-interval 1000
)

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.jsonl
   --val-prompt-data /fast/project/HFMI_SynergyUnit/tab_model/huggingface/nemo_hf_part_jsonl_4k_tokens_validation.jsonl
   --val-interval 1000
   --val-steps 100
   --input-key messages
   --apply-chat-template
   --rollout-shuffle
   --num-rollout 10000
   --rollout-batch-size 16
   --global-batch-size 16

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

LORA_ARGS=(
    --lora-rank 8
    --lora-alpha 16
    --target-modules q_proj,v_proj
)

OPTIMIZER_ARGS=(
   --optimizer adam
   --lr 1e-4
   --lr-decay-style WSD
   --lr-wsd-decay-style linear
   --lr-warmup-iters 500
   --lr-decay-iters 10000
   --lr-wsd-decay-iters 2000
   --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[@]}
[tab] requirements.txt (pip-requirements)
accelerate
blobfile
datasets
httpx[http2]
mcp[cli]
memray  # needed for debugging (but is lightweight), we can put it to dev mode when using pyproject.toml
omegaconf
peft
pillow
pylatexenc
pyyaml
qwen_vl_utils # for VLM
ray[default]
ring_flash_attn
sglang-router>=0.2.3
tensorboard
transformers
wandb
[content] vscode.git.Git (log)
2025-12-31 18:33:19.351 [info] > git ls-files --stage -- train_sft.py [9ms]
2025-12-31 18:33:19.356 [info] > git ls-files --stage -- scripts/run-sft-torchrun.sh [18ms]
2025-12-31 18:33:19.360 [info] > git ls-tree -l 1c50a119a28bd2331aaa72cb6634db29a9cdccdd -- train_sft.py [27ms]
2025-12-31 18:33:19.361 [info] > git ls-tree -l f0470e08f3dc27d0e1dc48884d46af25823e482e -- tests/models/peft/test_lora.py [31ms]
2025-12-31 18:33:19.362 [info] > git ls-tree -l 1c50a119a28bd2331aaa72cb6634db29a9cdccdd -- miles/utils/arguments.py [35ms]
2025-12-31 18:33:19.363 [info] > git ls-tree -l f0470e08f3dc27d0e1dc48884d46af25823e482e -- miles/models/peft/lora.py [39ms]
2025-12-31 18:33:19.364 [info] > git ls-tree -l f0470e08f3dc27d0e1dc48884d46af25823e482e -- miles/models/peft/arguments.py [44ms]
2025-12-31 18:33:19.365 [info] > git ls-tree -l f0470e08f3dc27d0e1dc48884d46af25823e482e -- miles/models/peft/__init__.py [48ms]
2025-12-31 18:33:19.366 [info] > git ls-tree -l f0470e08f3dc27d0e1dc48884d46af25823e482e -- miles/backends/fsdp_utils/checkpoint.py [53ms]
2025-12-31 18:33:19.367 [info] > git ls-tree -l 1c50a119a28bd2331aaa72cb6634db29a9cdccdd -- miles/backends/fsdp_utils/checkpoint.py [59ms]
2025-12-31 18:33:19.368 [info] > git ls-tree -l f0470e08f3dc27d0e1dc48884d46af25823e482e -- miles/backends/fsdp_utils/actor.py [64ms]
2025-12-31 18:33:19.369 [info] > git ls-tree -l 1c50a119a28bd2331aaa72cb6634db29a9cdccdd -- miles/backends/fsdp_utils/actor.py [68ms]
2025-12-31 18:33:19.369 [info] > git ls-tree -l f0470e08f3dc27d0e1dc48884d46af25823e482e -- train_sft.py [72ms]
2025-12-31 18:33:19.370 [info] > git ls-tree -l f0470e08f3dc27d0e1dc48884d46af25823e482e -- scripts/run-sft-torchrun.sh [76ms]
2025-12-31 18:33:19.371 [info] > git ls-tree -l f0470e08f3dc27d0e1dc48884d46af25823e482e -- miles/utils/arguments.py [80ms]
2025-12-31 18:33:19.373 [info] > git cat-file -s 70aa7888dc7a3ce8ff3230d8e5175d848eff9739 [13ms]
2025-12-31 18:33:19.374 [info] > git cat-file -s a2a1f981cf91f14828461ea9f9809ddc3eac7f33 [18ms]
2025-12-31 18:33:19.375 [info] > git ls-files --stage -- miles/backends/fsdp_utils/lora_utils.py [24ms]
2025-12-31 18:33:19.379 [info] > git ls-files --stage -- miles/backends/fsdp_utils/checkpoint.py [33ms]
2025-12-31 18:33:19.391 [info] > git cat-file -s 0df435e10fb2767bf21ef2d42527167e89e0beb1 [12ms]
2025-12-31 18:33:19.396 [info] > git cat-file -s 5287e047ea7159535668995f8a6c25a2b95bb222 [13ms]
[content] vscode.git.Git (log)
2025-12-31 18:33:19.464 [info] > git show --textconv 1c50a119a28bd2331aaa72cb6634db29a9cdccdd:train_sft.py [49ms]
2025-12-31 18:33:19.465 [info] > git show --textconv f0470e08f3dc27d0e1dc48884d46af25823e482e:tests/models/peft/test_lora.py [46ms]
2025-12-31 18:33:19.466 [info] > git show --textconv 1c50a119a28bd2331aaa72cb6634db29a9cdccdd:miles/utils/arguments.py [44ms]
2025-12-31 18:33:19.467 [info] > git show --textconv f0470e08f3dc27d0e1dc48884d46af25823e482e:miles/models/peft/lora.py [42ms]
2025-12-31 18:33:19.468 [info] > git show --textconv f0470e08f3dc27d0e1dc48884d46af25823e482e:miles/models/peft/arguments.py [39ms]
2025-12-31 18:33:19.469 [info] > git show --textconv f0470e08f3dc27d0e1dc48884d46af25823e482e:miles/models/peft/__init__.py [37ms]
2025-12-31 18:33:19.470 [info] > git show --textconv f0470e08f3dc27d0e1dc48884d46af25823e482e:miles/backends/fsdp_utils/checkpoint.py [35ms]
2025-12-31 18:33:19.471 [info] > git show --textconv 1c50a119a28bd2331aaa72cb6634db29a9cdccdd:miles/backends/fsdp_utils/checkpoint.py [32ms]
[content] vscode.git.Git (log)
2025-12-31 18:33:19.483 [info] > git show --textconv f0470e08f3dc27d0e1dc48884d46af25823e482e:miles/backends/fsdp_utils/actor.py [40ms]
2025-12-31 18:33:19.484 [info] > git show --textconv 1c50a119a28bd2331aaa72cb6634db29a9cdccdd:miles/backends/fsdp_utils/actor.py [36ms]
2025-12-31 18:33:19.485 [info] > git show --textconv f0470e08f3dc27d0e1dc48884d46af25823e482e:train_sft.py [32ms]
2025-12-31 18:33:19.486 [info] > git show --textconv f0470e08f3dc27d0e1dc48884d46af25823e482e:scripts/run-sft-torchrun.sh [28ms]
2025-12-31 18:33:19.492 [info] > git show --textconv f0470e08f3dc27d0e1dc48884d46af25823e482e:miles/utils/arguments.py [30ms]
2025-12-31 18:33:19.497 [info] > git show --textconv :scripts/run-sft-torchrun.sh [19ms]
2025-12-31 18:33:19.500 [info] > git show --textconv :train_sft.py [18ms]
2025-12-31 18:33:19.508 [info] > git show --textconv :miles/backends/fsdp_utils/lora_utils.py [17ms]
2025-12-31 18:33:19.513 [info] > git show --textconv :miles/backends/fsdp_utils/checkpoint.py [16ms]
[content] vscode.git.Git (log)
2025-12-31 18:33:19.755 [info] > git fetch [1733ms]
2025-12-31 18:33:19.766 [info] > git config --get commit.template [5ms]
2025-12-31 18:33:19.781 [info] > git for-each-ref --format=%(refname)%00%(upstream:short)%00%(objectname)%00%(upstream:track)%00%(upstream:remotename)%00%(upstream:remoteref) refs/heads/sync-lora-impl-with-upstream refs/remotes/sync-lora-impl-with-upstream [16ms]
[content] vscode.git.Git (log)
2025-12-31 18:33:20.940 [info] > git ls-files --stage -- train_sft.py [9ms]
2025-12-31 18:33:20.945 [info] > git ls-files --stage -- scripts/run-sft-torchrun.sh [19ms]
2025-12-31 18:33:20.949 [info] > git ls-tree -l 1c50a119a28bd2331aaa72cb6634db29a9cdccdd -- train_sft.py [27ms]
2025-12-31 18:33:20.951 [info] > git ls-tree -l f0470e08f3dc27d0e1dc48884d46af25823e482e -- tests/models/peft/test_lora.py [33ms]
2025-12-31 18:33:20.951 [info] > git ls-tree -l 1c50a119a28bd2331aaa72cb6634db29a9cdccdd -- miles/utils/arguments.py [36ms]
2025-12-31 18:33:20.952 [info] > git ls-tree -l f0470e08f3dc27d0e1dc48884d46af25823e482e -- miles/models/peft/lora.py [41ms]
2025-12-31 18:33:20.953 [info] > git ls-tree -l f0470e08f3dc27d0e1dc48884d46af25823e482e -- miles/models/peft/arguments.py [45ms]
2025-12-31 18:33:20.954 [info] > git ls-tree -l f0470e08f3dc27d0e1dc48884d46af25823e482e -- miles/models/peft/__init__.py [50ms]
2025-12-31 18:33:20.955 [info] > git ls-tree -l f0470e08f3dc27d0e1dc48884d46af25823e482e -- miles/backends/fsdp_utils/checkpoint.py [56ms]
2025-12-31 18:33:20.956 [info] > git ls-tree -l 1c50a119a28bd2331aaa72cb6634db29a9cdccdd -- miles/backends/fsdp_utils/checkpoint.py [61ms]
2025-12-31 18:33:20.957 [info] > git ls-tree -l f0470e08f3dc27d0e1dc48884d46af25823e482e -- miles/backends/fsdp_utils/actor.py [65ms]
2025-12-31 18:33:20.958 [info] > git ls-tree -l 1c50a119a28bd2331aaa72cb6634db29a9cdccdd -- miles/backends/fsdp_utils/actor.py [70ms]
2025-12-31 18:33:20.959 [info] > git ls-tree -l f0470e08f3dc27d0e1dc48884d46af25823e482e -- train_sft.py [74ms]
2025-12-31 18:33:20.960 [info] > git ls-tree -l f0470e08f3dc27d0e1dc48884d46af25823e482e -- scripts/run-sft-torchrun.sh [79ms]
2025-12-31 18:33:20.961 [info] > git ls-tree -l f0470e08f3dc27d0e1dc48884d46af25823e482e -- miles/utils/arguments.py [83ms]
2025-12-31 18:33:20.962 [info] > git cat-file -s 70aa7888dc7a3ce8ff3230d8e5175d848eff9739 [13ms]
2025-12-31 18:33:20.963 [info] > git cat-file -s a2a1f981cf91f14828461ea9f9809ddc3eac7f33 [18ms]
2025-12-31 18:33:20.963 [info] > git ls-files --stage -- miles/backends/fsdp_utils/lora_utils.py [24ms]
2025-12-31 18:33:20.968 [info] > git ls-files --stage -- miles/backends/fsdp_utils/checkpoint.py [32ms]
2025-12-31 18:33:20.980 [info] > git cat-file -s 0df435e10fb2767bf21ef2d42527167e89e0beb1 [12ms]
2025-12-31 18:33:20.984 [info] > git cat-file -s 5287e047ea7159535668995f8a6c25a2b95bb222 [12ms]
[content] vscode.git.Git (log)
2025-12-31 18:33:21.051 [info] > git show --textconv 1c50a119a28bd2331aaa72cb6634db29a9cdccdd:train_sft.py [48ms]
2025-12-31 18:33:21.052 [info] > git show --textconv f0470e08f3dc27d0e1dc48884d46af25823e482e:tests/models/peft/test_lora.py [45ms]
2025-12-31 18:33:21.053 [info] > git show --textconv 1c50a119a28bd2331aaa72cb6634db29a9cdccdd:miles/utils/arguments.py [43ms]
2025-12-31 18:33:21.054 [info] > git show --textconv f0470e08f3dc27d0e1dc48884d46af25823e482e:miles/models/peft/lora.py [40ms]
2025-12-31 18:33:21.055 [info] > git show --textconv f0470e08f3dc27d0e1dc48884d46af25823e482e:miles/models/peft/arguments.py [37ms]
2025-12-31 18:33:21.056 [info] > git show --textconv f0470e08f3dc27d0e1dc48884d46af25823e482e:miles/models/peft/__init__.py [35ms]
2025-12-31 18:33:21.057 [info] > git show --textconv f0470e08f3dc27d0e1dc48884d46af25823e482e:miles/backends/fsdp_utils/checkpoint.py [33ms]
2025-12-31 18:33:21.058 [info] > git show --textconv 1c50a119a28bd2331aaa72cb6634db29a9cdccdd:miles/backends/fsdp_utils/checkpoint.py [30ms]
2025-12-31 18:33:21.059 [info] > git show --textconv f0470e08f3dc27d0e1dc48884d46af25823e482e:miles/backends/fsdp_utils/actor.py [28ms]
2025-12-31 18:33:21.060 [info] > git show --textconv 1c50a119a28bd2331aaa72cb6634db29a9cdccdd:miles/backends/fsdp_utils/actor.py [25ms]
2025-12-31 18:33:21.062 [info] > git show --textconv f0470e08f3dc27d0e1dc48884d46af25823e482e:train_sft.py [22ms]
2025-12-31 18:33:21.072 [info] > git show --textconv f0470e08f3dc27d0e1dc48884d46af25823e482e:scripts/run-sft-torchrun.sh [26ms]
2025-12-31 18:33:21.078 [info] > git show --textconv f0470e08f3dc27d0e1dc48884d46af25823e482e:miles/utils/arguments.py [27ms]
2025-12-31 18:33:21.084 [info] > git show --textconv :scripts/run-sft-torchrun.sh [17ms]
2025-12-31 18:33:21.088 [info] > git show --textconv :train_sft.py [16ms]
2025-12-31 18:33:21.095 [info] > git show --textconv :miles/backends/fsdp_utils/lora_utils.py [17ms]
2025-12-31 18:33:21.100 [info] > git show --textconv :miles/backends/fsdp_utils/checkpoint.py [17ms]
[content] vscode.git.Git (log)
2025-12-31 18:33:21.883 [info] > git config --get commit.template [5ms]
2025-12-31 18:33:21.898 [info] > git for-each-ref --format=%(refname)%00%(upstream:short)%00%(objectname)%00%(upstream:track)%00%(upstream:remotename)%00%(upstream:remoteref) refs/heads/sync-lora-impl-with-upstream refs/remotes/sync-lora-impl-with-upstream [17ms]