p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] scripts/models/kimi-k2.sh (shellscript)
NLAYERS=61
FIRST_K_DENSE_REPLACE=1

arr=()
for ((i=0; i<NLAYERS; i++)); do
  if (( i < FIRST_K_DENSE_REPLACE )); then
    arr+=(0)
  else
    arr+=(1)
  fi
done

printf -v MOE_LAYER_FREQ "[%s]" "$(IFS=', '; echo "${arr[*]}")"

# kimi-k2
MODEL_ARGS=(
    --disable-bias-linear
    --num-layers 61
    --hidden-size 7168
    --ffn-hidden-size 18432
    --num-attention-heads 64
    --kv-channels 64
    --normalization RMSNorm
    --position-embedding-type rope
    --norm-epsilon 1e-6
    --swiglu
    --untie-embeddings-and-output-weights
    --vocab-size 163840
    
    --multi-latent-attention
    --q-lora-rank 1536
    --kv-lora-rank 512
    --qk-head-dim 128
    --qk-pos-emb-head-dim 64
    --v-head-dim 128
    --qk-layernorm
    --rotary-scaling-factor 32.0
    --rotary-base 50000
    --mscale 1.0
    --mscale-all-dim 1.0
    --attention-softmax-in-fp32
    --no-rope-fusion

    # moe
    --num-experts 384
    --moe-layer-freq $MOE_LAYER_FREQ
    --moe-ffn-hidden-size 2048
    --moe-router-topk 8
    --moe-shared-expert-intermediate-size 2048
    --moe-router-pre-softmax
    --moe-router-score-function sigmoid
    --moe-router-enable-expert-bias
    --moe-router-load-balancing-type seq_aux_loss
    --moe-token-dispatcher-type alltoall
    --moe-aux-loss-coeff 0
    --moe-router-bias-update-rate 0
    --moe-router-group-topk 1
    --moe-router-num-groups 1
    --moe-grouped-gemm
    --moe-router-topk-scaling-factor 2.827
    --moe-router-dtype fp32
    --moe-permute-fusion
)
[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_20251222_220632"
LOAD_PATH="/fast/project/HFMI_SynergyUnit/tab_model/huggingface/shared_data/${RUN_ID}/checkpoints"
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 200
)

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-4
   --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
   --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] miles/backends/fsdp_utils/actor.py (python)
            if args.use_lora:
                lora_config = LoRAConfig(
                    lora_rank=args.lora_rank,
                    lora_alpha=args.lora_alpha,
                    lora_dropout=args.lora_dropout,
                    target_modules=args.lora_target_modules,
                )
                model = apply_lora(model, lora_config)
                logger.info(f"[Rank {dist.get_rank()}] Applied LoRA: {lora_config}")
[content] miles/backends/fsdp_utils/actor.py (python)
            trainable_params = [p for p in self.model.parameters() if p.requires_grad]
            
            if 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 * trainable_count / total_params:.2f}%)"
                )
[content] miles/backends/sglang_utils/sglang_engine.py (python)
        
        # Try to find LoRA adapter files in checkpoint directory
        lora_paths = _find_lora_checkpoint_path(args)
        if lora_paths:
            kwargs["lora_paths"] = lora_paths
            logger.info(f"[Rank {rank}] Found LoRA adapter at {lora_paths}, will load at SGLang initialization")
[content] miles/backends/sglang_utils/sglang_engine.py (python)
def _find_lora_checkpoint_path(args) -> list[str] | None:
    """Find LoRA adapter files in checkpoint directory if available.
    
    Returns a list of paths to checkpoint directories containing LoRA adapters,
    or None if no checkpoint is found or adapter files don't exist.
    """
    load_root = getattr(args, "load", None)
    if load_root is None:
        return None
    
    root_path = Path(load_root).expanduser()
    if not root_path.exists():
        return None
    
    # Find the latest checkpoint (same logic as checkpoint.load)
    target_step = getattr(args, "ckpt_step", None)
    if target_step is None:
        tracker_file = root_path / "latest_checkpointed_iteration.txt"
        if not tracker_file.exists():
            return None
        tracker_text = tracker_file.read_text().strip()
        target_step = int(tracker_text)
    
    checkpoint_dir = root_path / f"iter_{target_step:07d}"
    
    # Check if LoRA adapter files exist
    adapter_config = checkpoint_dir / "adapter_config.json"
    adapter_model = checkpoint_dir / "adapter_model.safetensors"
    
    if adapter_config.exists() and adapter_model.exists():
        return [str(checkpoint_dir)]
    
    return None
[tab] scripts/run-qwen3-4B-fsdp.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 PYTHONBUFFERED=16
export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
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)"



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="/root/shared_data/${RUN_ID}/checkpoints"

CKPT_ARGS=(
   --hf-checkpoint /root/Qwen3-4B
   --load /root/Qwen3-4B
   --ref-load /root/Qwen3-4B
)

ROLLOUT_ARGS=(
   --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl
   --input-key prompt
   --label-key label
   --apply-chat-template
   --rollout-shuffle
   --balance-data
   --rm-type deepscaler
   --num-rollout 100
   --rollout-batch-size 8
   --n-samples-per-prompt 8
   --rollout-max-response-len 4096
   --rollout-temperature 0.8
   --global-batch-size 64
)

GRPO_ARGS=(
   --use-kl-loss
   --advantage-estimator grpo
   --kl-loss-coef 0.00
   --kl-loss-type low_var_kl
   --kl-coef 0.00
   --entropy-coef 0.00
   --eps-clip 0.2
   --eps-clip-high 0.28
)

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

WANDB_ARGS=(
   --use-wandb
   --wandb-project miles-dev-mcore-fsdp
   --wandb-group qwen3-4B-fsdp-1130-ref
   --wandb-key ${WANDB_API_KEY}
)

SGLANG_ARGS=(
   --rollout-num-gpus-per-engine 1
   --sglang-mem-fraction-static 0.75
   --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"}'
)

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

MISC_ARGS=(
   --actor-num-nodes 1
   --actor-num-gpus-per-node 8
   --colocate
   --use-fault-tolerance
   --dump-details /root/shared_data/qwen3-4B-fsdp-1116-noref/dump_details
   # --fsdp-cpu-offload
)

# launch the master node of ray in container - 8 GPUs for training
export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats


RUNTIME_ENV_JSON="{
  \"env_vars\": {
    \"PYTHONPATH\": \"/root/Megatron-LM/:${SCRIPT_DIR}\",
    \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\"
  }
}"


ray job submit --address="http://127.0.0.1:8265" \
   --runtime-env-json="${RUNTIME_ENV_JSON}" \
   -- python3 train.py \
   ${CKPT_ARGS[@]} \
   ${ROLLOUT_ARGS[@]} \
   ${OPTIMIZER_ARGS[@]} \
   ${GRPO_ARGS[@]} \
   ${WANDB_ARGS[@]} \
   ${SGLANG_ARGS[@]} \
   ${TRAIN_BACKEND_ARGS[@]} \
   ${PERF_ARGS[@]} \
   ${MISC_ARGS[@]}
[tab] /home/franz.srambical/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
[content] miles/backends/fsdp_utils/update_weight_utils.py (python)
            # Skip syncing LoRA weights entirely when loaded from files at SGLang init
            # because of name mismatch: training uses q_proj/v_proj but SGLang expects qkv_proj
            # TODO: Implement weight name transformation to map q_proj/v_proj -> qkv_proj
            if self.args.load:
                logger.warning(
                    f"[Rank {dist.get_rank()}] LoRA mode: skipping LoRA weight sync "
                    f"(loaded from files at SGLang init; name mismatch: q_proj/v_proj vs qkv_proj). "
                    f"This means SGLang will use the initial LoRA weights and won't receive updates during training."
                )
                return
[content] miles/backends/fsdp_utils/update_weight_utils.py (python)
            logger.info(
                f"[Rank {dist.get_rank()}] LoRA mode: syncing {len(state_dict)} LoRA weight tensors to SGLang. "
                f"Note: This may fail if there's a name mismatch (q_proj/v_proj vs qkv_proj)"
            )
[content] miles/backends/fsdp_utils/update_weight_utils.py (python)
            # Transform q_proj/v_proj LoRA weights to qkv_proj format for SGLang
            state_dict = _transform_lora_weights_for_sglang(state_dict)
            logger.info(
                f"[Rank {dist.get_rank()}] LoRA mode: syncing {len(state_dict)} LoRA weight tensors to SGLang "
                f"(transformed from q_proj/v_proj to qkv_proj format)"
            )
[content] miles/backends/fsdp_utils/update_weight_utils.py (python)
            
            # Transform LoRA weight names after converting to local tensor
            if self.args.use_lora and ("q_proj" in name or "v_proj" in name):
                # This will be handled in batch after collecting all weights
                pass
            
            bucket.append((name, param))
            bucket_size += param_size
        
        # Transform LoRA weights after collecting all of them and converting to local tensors
        if self.args.use_lora and bucket:
            # Extract names and params, transform, then rebuild bucket
            names_and_params = [(name, param) for name, param in bucket]
            state_dict_for_transform = {name: param for name, param in names_and_params}
            transformed_state_dict = _transform_lora_weights_for_sglang(state_dict_for_transform)
            
            # Rebuild bucket with transformed weights
            bucket = []
            bucket_size = 0
            for name, param in transformed_state_dict.items():
                param_size = param.numel() * param.element_size()
                bucket.append((name, param))
                bucket_size += param_size
            
            logger.info(
                f"[Rank {dist.get_rank()}] LoRA mode: transformed to {len(transformed_state_dict)} qkv_proj weights"
            )
[content] miles/backends/fsdp_utils/update_weight_utils.py (python)
            logger.info(
                f"[Rank {dist.get_rank()}] LoRA mode: found {len(state_dict)} LoRA weight tensors, "
                f"transforming q_proj/v_proj to qkv_proj format for SGLang"
            )
[content] miles/backends/fsdp_utils/update_weight_utils.py (python)
            param = param.cuda()
            if isinstance(param, DTensor):
                # Convert DTensor to local tensor (synchronously for transformation)
                param = param.redistribute(
                    placements=[Replicate()] * param.device_mesh.ndim,
                    async_op=False,
                ).to_local()
            local_state_dict[name] = param
        
        # Transform q_proj/v_proj to qkv_proj format after converting to local tensors
        if self.args.use_lora:
            local_state_dict = _transform_lora_weights_for_sglang(local_state_dict)
            logger.info(
                f"[Rank {dist.get_rank()}] LoRA mode: transformed to {len(local_state_dict)} qkv_proj weights"
            )
        
        # Now process the transformed weights into buckets
        for name, param in local_state_dict.items():
[content] miles/backends/fsdp_utils/update_weight_utils.py (python)
            param = param.cuda()
            if isinstance(param, DTensor):
                # async version of param.full_tensor
                param = param.redistribute(
                    placements=[Replicate()] * param.device_mesh.ndim,
                    async_op=True,
                ).to_local()
[content] miles/backends/fsdp_utils/update_weight_utils.py (python)
            if isinstance(param, DTensor):
                # Convert DTensor to local tensor first (must be on CUDA device for redistribute to work)
                # Move to CUDA first, then redistribute
                if param.device.type == "cpu":
                    # If on CPU, we need to move to CUDA first, but DTensor.cuda() doesn't work
                    # So we need to redistribute on the current device, then move
                    # Actually, we should ensure the DTensor is already on CUDA from the model
                    # But if it's not, we need to handle it differently
                    logger.warning(
                        f"[Rank {dist.get_rank()}] LoRA weight {name} is a DTensor on CPU, "
                        f"this may cause issues. Device mesh: {param.device_mesh}"
                    )
                # Redistribute to Replicate and get local tensor (synchronously for transformation)
                param = param.redistribute(
                    placements=[Replicate()] * param.device_mesh.ndim,
                    async_op=False,
                ).to_local()
            # Now move to CUDA if it's a regular tensor
            if not isinstance(param, DTensor):
                param = param.cuda()
            local_state_dict[name] = param
        
        # Transform q_proj/v_proj to qkv_proj format after converting to local tensors
        if self.args.use_lora:
            local_state_dict = _transform_lora_weights_for_sglang(local_state_dict)
            logger.info(
                f"[Rank {dist.get_rank()}] LoRA mode: transformed to {len(local_state_dict)} qkv_proj weights"
            )
        
        # Now process the transformed weights into buckets
        for name, param in local_state_dict.items():
[content] miles/backends/fsdp_utils/update_weight_utils.py (python)
        # Collect all LoRA weights and convert DTensors to local tensors first
        # IMPORTANT: Convert DTensor to local tensor BEFORE moving to CUDA, because
        # DTensor.cuda() doesn't work properly and redistribute() on CPU DTensors fails
        local_state_dict = {}
[content] miles/backends/fsdp_utils/update_weight_utils.py (python)
            if isinstance(param, DTensor):
                # Check if the DTensor's underlying device is CPU (which would cause the error)
                # The device property of DTensor might not reflect the actual data location
                # We need to redistribute to Replicate first, which should work if the model is on CUDA
                try:
                    # Redistribute to Replicate and get local tensor (synchronously for transformation)
                    param = param.redistribute(
                        placements=[Replicate()] * param.device_mesh.ndim,
                        async_op=False,
                    ).to_local()
                    # Now move to CUDA
                    param = param.cuda()
                except RuntimeError as e:
                    if "No backend type associated with device type cpu" in str(e):
                        # If redistribute fails because of CPU backend, try to get the local tensor directly
                        # This might work if the DTensor is already sharded
                        logger.error(
                            f"[Rank {dist.get_rank()}] Failed to redistribute DTensor {name}: {e}. "
                            f"Trying alternative approach..."
                        )
                        # Try using full_tensor() if available, or skip this weight
                        try:
                            param = param.full_tensor().cuda()
                        except:
                            logger.error(f"[Rank {dist.get_rank()}] Cannot convert DTensor {name}, skipping")
                            continue
                    else:
                        raise
            else:
                # Regular tensor, just move to CUDA
                param = param.cuda()
            local_state_dict[name] = param
        
        # Transform q_proj/v_proj to qkv_proj format after converting to local tensors
        if self.args.use_lora:
            local_state_dict = _transform_lora_weights_for_sglang(local_state_dict)
            logger.info(
                f"[Rank {dist.get_rank()}] LoRA mode: transformed to {len(local_state_dict)} qkv_proj weights"
            )
        
        # Now process the transformed weights into buckets
        for name, param in local_state_dict.items():
[content] miles/backends/fsdp_utils/update_weight_utils.py (python)
        # Collect all LoRA weights and convert DTensors to local tensors first
        # IMPORTANT: The issue is that DTensor.redistribute() on CPU tensors tries to use CPU backend
        # which doesn't exist. We need to ensure DTensors are handled correctly.
        # FSDP DTensors should already be on CUDA, but we need to redistribute them properly.
        local_state_dict = {}
[content] miles/backends/fsdp_utils/update_weight_utils.py (python)
            if isinstance(param, DTensor):
                # For single process or CPU DTensors, use full_tensor() to avoid backend issues
                if dist.get_world_size() == 1 or param.device.type == "cpu":
                    param = param.full_tensor()
                    if param.device.type == "cpu":
                        param = param.cuda()
                else:
                    # For multi-process CUDA DTensors, redistribute to Replicate
                    # Redistribute to Replicate and get local tensor (synchronously for transformation)
                    param = param.redistribute(
                        placements=[Replicate()] * param.device_mesh.ndim,
                        async_op=False,
                    ).to_local()
                    # Ensure on CUDA
                    if param.device.type != "cuda":
                        param = param.cuda()
            else:
                # Regular tensor, just move to CUDA
                param = param.cuda()
            local_state_dict[name] = param
        
        # Transform q_proj/v_proj to qkv_proj format after converting to local tensors
        if self.args.use_lora:
            local_state_dict = _transform_lora_weights_for_sglang(local_state_dict)
            logger.info(
                f"[Rank {dist.get_rank()}] LoRA mode: transformed to {len(local_state_dict)} qkv_proj weights"
            )
        
        # Now process the transformed weights into buckets
        for name, param in local_state_dict.items():
[content] miles/backends/fsdp_utils/update_weight_utils.py (python)
        # Collect all LoRA weights and convert DTensors to local tensors first
        # IMPORTANT: DTensor.redistribute() requires a distributed backend. If the DTensor is on CPU,
        # it tries to use CPU/Gloo backend which isn't initialized (only CUDA/NCCL is).
        # Solution: Use full_tensor() for world_size=1 or when on CPU, otherwise redistribute.
        local_state_dict = {}
[content] miles/backends/fsdp_utils/update_weight_utils.py (python)
        # Use get_model_state_dict() instead of model.state_dict() to properly handle FSDP CPU offload
        # This ensures DTensors are materialized on CUDA even if CPUOffloadPolicy is enabled
        from torch.distributed.checkpoint.state_dict import get_model_state_dict, StateDictOptions
        
        # Get state dict with proper FSDP handling (materializes offloaded params to CUDA)
        options = StateDictOptions(full_state_dict=True, cpu_offload=False)
        full_state_dict, _ = get_model_state_dict(self.model, optimizers=[], options=options)
[content] miles/backends/fsdp_utils/update_weight_utils.py (python)
        # Use get_model_state_dict() instead of model.state_dict() to properly handle FSDP CPU offload
        # This ensures DTensors are materialized on CUDA even if CPUOffloadPolicy is enabled
        from torch.distributed.checkpoint.state_dict import get_model_state_dict, StateDictOptions
        
        # Get state dict with proper FSDP handling (materializes offloaded params to CUDA)
        options = StateDictOptions(full_state_dict=True, cpu_offload=False)
        full_state_dict = get_model_state_dict(self.model, options=options)
[content] miles/backends/sglang_utils/sglang_engine.py (python)
        lora_checkpoint_dir = _find_lora_checkpoint_path(args)
        logger.info(f"[Rank {rank}] Loading LoRA adapter from {lora_checkpoint_dir}")
        kwargs["lora_paths"] = [f"actor_lora={lora_checkpoint_dir}"]
[content] miles/backends/fsdp_utils/update_weight_utils.py (python)
            state_dict = {k: v for k, v in state_dict.items() if "lora_" not in k}
            logger.info(f"[Rank {dist.get_rank()}] LoRA mode: skipping LoRA weights (base model weights only). "
                       f"LoRA adapter updates should use unload_lora_adapter/load_lora_adapter endpoints.")
[content] miles/backends/fsdp_utils/update_weight_utils.py (python)
        if self.args.use_lora:
            state_dict = {k: v for k, v in state_dict.items() if "lora_" in k}
            logger.info(f"[Rank {dist.get_rank()}] LoRA mode: syncing only LoRA weights to SGLang")
[content] miles/backends/fsdp_utils/update_weight_utils.py (python)
    def update_lora_adapter(self) -> None:
        """Update LoRA adapter using SGLang's unload_lora_adapter and load_lora_adapter endpoints.

        Since update_weights_from_tensor doesn't support LoRA weights, we need to:
        1. Save current LoRA weights to temporary adapter files
        2. Unload current LoRA adapter from SGLang
        3. Load new LoRA adapter from temporary files
        """
        logger.warning(f"[Rank {dist.get_rank()}] LoRA adapter update not yet implemented. "
                      f"Skipping LoRA weight sync. Base model weights are already in SGLang.")
        # TODO: Implement LoRA adapter update logic:
        # 1. Extract LoRA weights from self.model.state_dict()
        # 2. Save to temporary adapter_config.json and adapter_model.safetensors
        # 3. Call unload_lora_adapter on SGLang engines
        # 4. Call load_lora_adapter on SGLang engines with new adapter files
[content] miles/backends/fsdp_utils/update_weight_utils.py (python)
    def _update_lora_weights(self) -> None:
        """Update LoRA weights by saving to disk and using SGLang's adapter endpoints."""
        logger.info(f"[Rank {dist.get_rank()}] LoRA mode: using load/unload adapter endpoints")
        
        # Get full state dict with LoRA weights on rank 0
        if dist.get_rank() == 0:
            options = StateDictOptions(full_state_dict=True, cpu_offload=True)
            full_state_dict = get_model_state_dict(self.model, options=options)
            lora_state_dict = {k: v for k, v in full_state_dict.items() if "lora_" in k}
            
            if not lora_state_dict:
                logger.warning("[Rank 0] No LoRA weights found in model state dict")
                return
            
            # Create temporary directory for LoRA adapter
            with tempfile.TemporaryDirectory() as temp_dir:
                adapter_dir = Path(temp_dir)
                
                # Save LoRA config
                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,
                )
                peft_config = lora_config.to_hf_peft_config()
                with open(adapter_dir / "adapter_config.json", "w") as f:
                    json.dump(peft_config, f, indent=2)
                
                # Save LoRA weights
                safetensors.torch.save_file(lora_state_dict, adapter_dir / "adapter_model.safetensors")
                logger.info(f"[Rank 0] Saved LoRA adapter to {adapter_dir}")
                
                # Unload and load adapter on all rollout engines
                # Note: We need to get rollout_engines, but they should be set in connect_rollout_engines
                # For now, we'll handle this in the subclasses that have access to engines
                if hasattr(self, 'rollout_engines'):
                    adapter_path = str(adapter_dir)
                    for engine in self.rollout_engines:
                        # Unload current adapter
                        ref_unload = engine.unload_lora_adapter.remote(adapter_name="actor_lora")
                        ray.get(ref_unload)
                        # Load new adapter
                        ref_load = engine.load_lora_adapter.remote(adapter_path=adapter_path, adapter_name="actor_lora")
                        ray.get(ref_load)
                    logger.info(f"[Rank 0] Updated LoRA adapter on all rollout engines")
                else:
                    logger.warning("[Rank 0] Rollout engines not connected, skipping LoRA update")
        
        dist.barrier()
[content] miles/backends/fsdp_utils/update_weight_utils.py (python)
        # When using LoRA, use SGLang's load/unload adapter endpoints instead of update_weights_from_tensor
        if self.args.use_lora:
            self._update_lora_weights()
            return
        
        # For non-LoRA weights, use the standard tensor update approach
[content] miles/backends/fsdp_utils/update_weight_utils.py (python)
    def _update_lora_weights(self) -> None:
        """Update LoRA weights by saving to disk and using SGLang's adapter endpoints."""
        logger.info(f"[Rank {dist.get_rank()}] LoRA mode: using load/unload adapter endpoints")
        
        # Get full state dict with LoRA weights on rank 0
        if dist.get_rank() == 0:
            options = StateDictOptions(full_state_dict=True, cpu_offload=True)
            full_state_dict = get_model_state_dict(self.model, options=options)
            lora_state_dict = {k: v for k, v in full_state_dict.items() if "lora_" in k}
            
            if not lora_state_dict:
                logger.warning("[Rank 0] No LoRA weights found in model state dict")
                return
            
            # Create temporary directory for LoRA adapter (persist until after load completes)
            temp_dir = tempfile.mkdtemp()
            adapter_dir = Path(temp_dir)
            
            try:
                # Save LoRA config
                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,
                )
                peft_config = lora_config.to_hf_peft_config()
                with open(adapter_dir / "adapter_config.json", "w") as f:
                    json.dump(peft_config, f, indent=2)
                
                # Save LoRA weights
                safetensors.torch.save_file(lora_state_dict, adapter_dir / "adapter_model.safetensors")
                logger.info(f"[Rank 0] Saved LoRA adapter to {adapter_dir}")
                
                # Unload and load adapter on all rollout engines
                # Note: We need to get rollout_engines, but they should be set in connect_rollout_engines
                # For now, we'll handle this in the subclasses that have access to engines
                if hasattr(self, 'rollout_engines'):
                    adapter_path = str(adapter_dir)
                    for engine in self.rollout_engines:
                        # Unload current adapter
                        ref_unload = engine.unload_lora_adapter.remote(adapter_name="actor_lora")
                        ray.get(ref_unload)
                        # Load new adapter
                        ref_load = engine.load_lora_adapter.remote(adapter_path=adapter_path, adapter_name="actor_lora")
                        ray.get(ref_load)
                    logger.info(f"[Rank 0] Updated LoRA adapter on all rollout engines")
                else:
                    logger.warning("[Rank 0] Rollout engines not connected, skipping LoRA update")
            finally:
                # Clean up temporary directory after adapter is loaded
                import shutil
                shutil.rmtree(temp_dir, ignore_errors=True)
        
        dist.barrier()
[content] miles/backends/fsdp_utils/update_weight_utils.py (python)
    def _update_lora_weights(self) -> None:
        """Update LoRA weights by saving to disk and using SGLang's adapter endpoints."""
        logger.info(f"[Rank {dist.get_rank()}] LoRA mode: using load/unload adapter endpoints")
        
        # Get full state dict with LoRA weights on rank 0
        if dist.get_rank() == 0:
            options = StateDictOptions(full_state_dict=True, cpu_offload=True)
            full_state_dict = get_model_state_dict(self.model, options=options)
            lora_state_dict = {k: v for k, v in full_state_dict.items() if "lora_" in k}
            
            if not lora_state_dict:
                logger.warning("[Rank 0] No LoRA weights found in model state dict")
                return
            
            # Create temporary directory for LoRA adapter (persist until after load completes)
            temp_dir = tempfile.mkdtemp()
            adapter_dir = Path(temp_dir)
            
            try:
                # Save LoRA config
                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,
                )
                peft_config = lora_config.to_hf_peft_config()
                with open(adapter_dir / "adapter_config.json", "w") as f:
                    json.dump(peft_config, f, indent=2)
                
                # Save LoRA weights (convert to regular dict for safetensors)
                lora_tensors = {k: v for k, v in lora_state_dict.items() if isinstance(v, torch.Tensor)}
                safetensors.torch.save_file(lora_tensors, adapter_dir / "adapter_model.safetensors")
                logger.info(f"[Rank 0] Saved LoRA adapter to {adapter_dir}")
                
                # Unload and load adapter on all rollout engines
                # Note: rollout_engines should be set in connect_rollout_engines by subclasses
                if hasattr(self, 'rollout_engines') and self.rollout_engines:
                    adapter_path = str(adapter_dir)
                    for engine in self.rollout_engines:
                        # Unload current adapter
                        ref_unload = engine.unload_lora_adapter.remote(adapter_name="actor_lora")
                        ray.get(ref_unload)
                        # Load new adapter
                        ref_load = engine.load_lora_adapter.remote(adapter_path=adapter_path, adapter_name="actor_lora")
                        ray.get(ref_load)
                    logger.info(f"[Rank 0] Updated LoRA adapter on all rollout engines")
                else:
                    logger.warning("[Rank 0] Rollout engines not connected, skipping LoRA update")
            finally:
                # Clean up temporary directory after adapter is loaded
                shutil.rmtree(temp_dir, ignore_errors=True)
        
        dist.barrier()
[content] miles/backends/fsdp_utils/update_weight_utils.py (python)
    def _update_lora_weights(self) -> None:
        """Update LoRA weights by saving to disk and using SGLang's adapter endpoints."""
        logger.info(f"[Rank {dist.get_rank()}] LoRA mode: using load/unload adapter endpoints")
        
        # Get full state dict with LoRA weights on rank 0
        if dist.get_rank() == 0:
            options = StateDictOptions(full_state_dict=True, cpu_offload=True)
            full_state_dict = get_model_state_dict(self.model, options=options)
            lora_state_dict = {k: v for k, v in full_state_dict.items() if "lora_" in k}
            
            if not lora_state_dict:
                logger.warning("[Rank 0] No LoRA weights found in model state dict")
                return
            
            # Create temporary directory for LoRA adapter (persist until after load completes)
            temp_dir = tempfile.mkdtemp()
            adapter_dir = Path(temp_dir)
            
            try:
                # Save LoRA config
                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,
                )
                peft_config = lora_config.to_hf_peft_config()
                with open(adapter_dir / "adapter_config.json", "w") as f:
                    json.dump(peft_config, f, indent=2)
                
                # Save LoRA weights (convert to regular dict for safetensors)
                lora_tensors = {k: v for k, v in lora_state_dict.items() if isinstance(v, torch.Tensor)}
                safetensors.torch.save_file(lora_tensors, adapter_dir / "adapter_model.safetensors")
                logger.info(f"[Rank 0] Saved LoRA adapter to {adapter_dir}")
                
                # Unload and load adapter on all rollout engines
                # Note: rollout_engines should be set in connect_rollout_engines by subclasses
                rollout_engines = getattr(self, 'rollout_engines', None)
                if rollout_engines:
                    adapter_path = str(adapter_dir)
                    for engine in rollout_engines:
                        # Unload current adapter
                        ref_unload = engine.unload_lora_adapter.remote(adapter_name="actor_lora")
                        ray.get(ref_unload)
                        # Load new adapter
                        ref_load = engine.load_lora_adapter.remote(adapter_path=adapter_path, adapter_name="actor_lora")
                        ray.get(ref_load)
                    logger.info(f"[Rank 0] Updated LoRA adapter on all rollout engines")
                else:
                    logger.warning("[Rank 0] Rollout engines not connected, skipping LoRA update")
            finally:
                # Clean up temporary directory after adapter is loaded
                shutil.rmtree(temp_dir, ignore_errors=True)
        
        dist.barrier()
[content] miles/backends/fsdp_utils/update_weight_utils.py (python)
    def _update_lora_weights(self) -> None:
        """Update LoRA weights by saving to disk and using SGLang's adapter endpoints."""
        logger.info(f"[Rank {dist.get_rank()}] LoRA mode: using load/unload adapter endpoints")
        
        # Get full state dict with LoRA weights on rank 0
        if dist.get_rank() == 0:
            options = StateDictOptions(full_state_dict=True, cpu_offload=True)
            full_state_dict = get_model_state_dict(self.model, options=options)
            lora_state_dict = {k: v for k, v in full_state_dict.items() if "lora_" in k}
            
            if not lora_state_dict:
                logger.warning("[Rank 0] No LoRA weights found in model state dict")
                return
            
            # Create temporary directory for LoRA adapter (persist until after load completes)
            temp_dir = tempfile.mkdtemp()
            adapter_dir = Path(temp_dir)
            
            try:
                # Save LoRA weights only (config is already set in SGLang from initial load)
                lora_tensors = {k: v for k, v in lora_state_dict.items() if isinstance(v, torch.Tensor)}
                safetensors.torch.save_file(lora_tensors, adapter_dir / "adapter_model.safetensors")
                logger.info(f"[Rank 0] Saved LoRA weights to {adapter_dir}")
                
                # Unload and load adapter on all rollout engines
                # Note: rollout_engines should be set in connect_rollout_engines by subclasses
                rollout_engines = getattr(self, 'rollout_engines', None)
                if rollout_engines:
                    adapter_path = str(adapter_dir)
                    for engine in rollout_engines:
                        # Unload current adapter
                        ref_unload = engine.unload_lora_adapter.remote(adapter_name="actor_lora")
                        ray.get(ref_unload)
                        # Load new adapter
                        ref_load = engine.load_lora_adapter.remote(adapter_path=adapter_path, adapter_name="actor_lora")
                        ray.get(ref_load)
                    logger.info(f"[Rank 0] Updated LoRA adapter on all rollout engines")
                else:
                    logger.warning("[Rank 0] Rollout engines not connected, skipping LoRA update")
            finally:
                # Clean up temporary directory after adapter is loaded
                shutil.rmtree(temp_dir, ignore_errors=True)
        
        dist.barrier()
[content] miles/backends/fsdp_utils/update_weight_utils.py (python)
    def _update_lora_weights(self) -> None:
        """Update LoRA weights by saving to disk and using SGLang's adapter endpoints."""
        logger.info(f"[Rank {dist.get_rank()}] LoRA mode: using load/unload adapter endpoints")
        
        # Get full state dict with LoRA weights on rank 0
        if dist.get_rank() == 0:
            options = StateDictOptions(full_state_dict=True, cpu_offload=True)
            full_state_dict = get_model_state_dict(self.model, options=options)
            lora_state_dict = {k: v for k, v in full_state_dict.items() if "lora_" in k}
            
            if not lora_state_dict:
                logger.warning("[Rank 0] No LoRA weights found in model state dict")
                return
            
            # Create temporary directory for LoRA adapter (persist until after load completes)
            temp_dir = tempfile.mkdtemp()
            adapter_dir = Path(temp_dir)
            
            try:
                # SGLang's load_lora_adapter endpoint requires a file path, so we must save weights to disk
                # We only save the weights file (not config) since the adapter config is already set in SGLang
                lora_tensors = {k: v for k, v in lora_state_dict.items() if isinstance(v, torch.Tensor)}
                safetensors.torch.save_file(lora_tensors, adapter_dir / "adapter_model.safetensors")
                logger.info(f"[Rank 0] Saved LoRA weights to {adapter_dir} for SGLang adapter update")
                
                # Unload and load adapter on all rollout engines
                # Note: rollout_engines should be set in connect_rollout_engines by subclasses
                rollout_engines = getattr(self, 'rollout_engines', None)
                if rollout_engines:
                    adapter_path = str(adapter_dir)
                    for engine in rollout_engines:
                        # Unload current adapter
                        ref_unload = engine.unload_lora_adapter.remote(adapter_name="actor_lora")
                        ray.get(ref_unload)
                        # Load new adapter
                        ref_load = engine.load_lora_adapter.remote(adapter_path=adapter_path, adapter_name="actor_lora")
                        ray.get(ref_load)
                    logger.info(f"[Rank 0] Updated LoRA adapter on all rollout engines")
                else:
                    logger.warning("[Rank 0] Rollout engines not connected, skipping LoRA update")
            finally:
                # Clean up temporary directory after adapter is loaded
                shutil.rmtree(temp_dir, ignore_errors=True)
        
        dist.barrier()
[tab] miles/utils/misc.py (python)
import importlib
import subprocess

import ray

from miles.utils.http_utils import is_port_available


def load_function(path):
    """
    Load a function from a module.
    :param path: The path to the function, e.g. "module.submodule.function".
    :return: The function object.
    """
    module_path, _, attr = path.rpartition(".")
    module = importlib.import_module(module_path)
    return getattr(module, attr)


class SingletonMeta(type):
    """
    A metaclass for creating singleton classes.
    """

    _instances = {}

    def __call__(cls, *args, **kwargs):
        if cls not in cls._instances:
            instance = super().__call__(*args, **kwargs)
            cls._instances[cls] = instance
        return cls._instances[cls]


def exec_command(cmd: str, capture_output: bool = False) -> str | None:
    print(f"EXEC: {cmd}", flush=True)

    try:
        result = subprocess.run(
            ["bash", "-c", cmd],
            shell=False,
            check=True,
            capture_output=capture_output,
            **(dict(text=True) if capture_output else {}),
        )
    except subprocess.CalledProcessError as e:
        if capture_output:
            print(f"{e.stdout=} {e.stderr=}")
        raise

    if capture_output:
        print(f"Captured stdout={result.stdout} stderr={result.stderr}")
        return result.stdout


def get_current_node_ip():
    address = ray._private.services.get_node_ip_address()
    # strip ipv6 address
    address = address.strip("[]")
    return address


def get_free_port(start_port=10000, consecutive=1):
    # find the port where port, port + 1, port + 2, ... port + consecutive - 1 are all available
    port = start_port
    while not all(is_port_available(port + i) for i in range(consecutive)):
        port += 1
    return port


def should_run_periodic_action(
    rollout_id: int,
    interval: int | None,
    num_rollout_per_epoch: int | None = None,
) -> bool:
    """
    Return True when a periodic action (eval/save/checkpoint) should run.

    Args:
        rollout_id: The current rollout index (0-based).
        interval: Desired cadence; disables checks when None.
        num_rollout_per_epoch: Optional epoch boundary to treat as a trigger.
    """
    if interval is None:
        return False

    step = rollout_id + 1
    return (step % interval == 0) or (num_rollout_per_epoch is not None and step % num_rollout_per_epoch == 0)