sft_utils

The supervised fine-tuning core: dataset loading and formatting, chain-of-thought prompt filling, LoRA and full-parameter trainer construction, and the temperature-based multi-task sampler.

medvision_bm.sft.sft_utils.is_main_process()

medvision_bm.sft.sft_utils.safe_print(*args, force=False, **kwargs)

Print only on main process unless force=True.

medvision_bm.sft.sft_utils.broadcast_int_from_main(value, src=0)

medvision_bm.sft.sft_utils.broadcast_object_from_main(value, src=0)

Broadcast any picklable object from the main process to all other processes.

Same rationale as broadcast_int_from_main(); used e.g. for a prepared-dataset path that only the main process can resolve. Returns value unchanged when not running distributed.

medvision_bm.sft.sft_utils.get_cgroup_limited_cpus()

medvision_bm.sft.sft_utils.get_cgroup_memory_percent()

Return (used_gib, limit_gib, percent) for the memory cgroup.

Reads cgroup v2 (memory.current / memory.max), falling back to cgroup v1 and finally to host-wide psutil numbers. On pods psutil reports the HOST total, so the cgroup files are the only trustworthy signal there.

medvision_bm.sft.sft_utils.get_image_info_for_medvision_dataset(doc)

Get the image modality and label name for a MedVision sample.

The sample’s taskType (defined in MedVision.py on the dataset repo) selects the dataset-specific preprocessing module, from which the image modality and the human-readable label name are looked up.

Parameters:

doc – A data sample from the MedVision dataset. Must contain taskType, dataset_name and taskID; label is present for all task types except the Biometrics-From-Landmarks (angle / distance) tasks.

Returns:

(image_modality, label_name), where label_name is None for tasks that have no label (the angle / distance tasks).

Return type:

tuple

Raises:

ValueError – If taskType is not one of the valid task types, or the dataset is not registered in DATASETS_NAME2PACKAGE.

medvision_bm.sft.sft_utils.normalize_ct_img(img, window_width, window_level)

Normalizes CT Hounsfield Units to [0, 255] based on W and L.

medvision_bm.sft.sft_utils.normalize_general_img(img)

Standard min-max normalization to [0, 255] for MR, PET, etc.

medvision_bm.sft.sft_utils.normalize_img(doc, img_2d)

Convert document to image with scale bar added.

medvision_bm.sft.sft_utils.img_proccessor_nii2png_save2disk(example, new_shape_hw=None)

medvision_bm.sft.sft_utils.img_proccessor_nii2png_save2dataset(example, new_shape_hw=None)

medvision_bm.sft.sft_utils.safe_concat_align_top_keys(datasets_list, fill_value=None)

Concatenate Hugging Face datasets even if they have different top-level keys. Missing columns are added and filled with fill_value.

medvision_bm.sft.sft_utils.safe_concat_align_dict_keys(datasets_list, dict_cols=None, fill_value=None)

Concatenate Hugging Face datasets even if dict columns have different keys.

Parameters:

• datasets_list (list[Dataset]) – list of datasets to concatenate

• dict_cols (list[str] | None) – names of columns containing dicts (if None, auto-detects)

• fill_value (any) – value used to fill missing keys (default: None)

Returns:

concatenated dataset

Return type:

Dataset

medvision_bm.sft.sft_utils.safe_concatenate_datasets(datasets_list)

medvision_bm.sft.sft_utils.group_train_test_split(dataset, group_column, test_size, seed=None, stratify_column=None)

Splits a HF Dataset into train and validation sets ensuring samples with the same value in ‘group_column’ are in the same split.

Parameters:

• dataset – The HF Dataset to split.

• group_column – The column name to group by (e.g., ‘image_file’).

• test_size – If float < 1.0, represents fraction of samples to aim for. If int >= 1, represents exact number of samples to aim for.

• seed – Random seed for shuffling.

• stratify_column – Optional column name to stratify by (e.g., ‘dataset_name’). A stratum is the set of volumes belonging to one unique value of this column (e.g., all volumes from ‘BraTS24’ form one stratum, all volumes from ‘AMOS22’ form another). When provided, volumes are interleaved round-robin across strata — one volume per stratum per round — before the greedy allocation loop runs. This guarantees every stratum contributes at least one volume to val before any stratum gets a second, preventing a few large-volume datasets from monopolising the val quota.

Returns:

DatasetDict containing ‘train’ and ‘validation’.

medvision_bm.sft.sft_utils.load_split_limit_dataset(tasks_list_json_path, limit_train_sample, limit_val_sample, num_workers_concat_datasets=4, tag_ds=None, download_mode='reuse_dataset_if_exists')

Load MedVision tasks, concatenate them, and split into train/validation.

Reads the task list from tasks_list_json_path and loads each task’s _Train split in parallel (falling back to single-threaded loading when any dataset was just downloaded, to avoid cache conflicts). The per-task datasets are concatenated in the deterministic JSON order (not arrival order) so the seeded shuffle and split downstream stay reproducible, then split into train and validation sets grouped by image_file (to prevent 3D-volume leakage) and stratified by dataset_name.

Parameters:

• tasks_list_json_path (str) – Path to the JSON file whose keys are the task names to load.

• limit_train_sample (int) – Cap on training samples after concatenation. Use a value < 0 for no limit or > 0 for a fixed cap; 0 is rejected.

• limit_val_sample (int) – Target size of the validation split; must be > 0.

• num_workers_concat_datasets (int) – Requested worker processes for parallel loading; clamped to the CPU count and task count. Defaults to 4.

• tag_ds (str) – Tag embedded in task names (<dataset_name>_<tag_ds>), used to recover the dataset name from each task. Required.

• download_mode (str) – Passed through to the dataset loader. Defaults to "reuse_dataset_if_exists".

Returns:

A dict with "train" and "validation" splits.

Return type:

DatasetDict

Raises:

• AssertionError – If limit_val_sample is not > 0, limit_train_sample is 0, tag_ds is None, or MedVision_DATA_DIR is unset.

• RuntimeError – If any task fails to load.

medvision_bm.sft.sft_utils.format_dataset(dataset, mapping_func, mapping_func_args, num_workers_format_dataset, writer_batch_size=1000)

Apply a formatting map function to a dataset with bounded parallelism.

Runs dataset.map(mapping_func, ...) to convert raw MedVision rows into the chat messages format expected by the SFT trainer. The number of worker processes is capped at the cgroup-limited CPU count.

Parameters:

• dataset – A HuggingFace Dataset or DatasetDict to format.

• mapping_func – The per-example formatting function passed to .map().

• mapping_func_args (dict) – Keyword arguments forwarded to mapping_func via fn_kwargs.

• num_workers_format_dataset (int) – Requested number of worker processes; clamped to the available CPU count.

• writer_batch_size (int) – Number of rows buffered before writing to the Arrow cache. Defaults to 1000.

Returns:

The formatted dataset with the mapping function applied.

medvision_bm.sft.sft_utils.clean_dataset(dataset, keys_to_keep)

Drop all columns from a dataset except a whitelist of keys.

Uses remove_columns — a schema-only operation — so no pass over the data is made, keeping the cached rows small before training.

Parameters:

• dataset – A HuggingFace Dataset or DatasetDict to prune.

• keys_to_keep (list[str]) – Column names to retain; all other columns are removed.

Returns:

The dataset containing only the whitelisted columns.

medvision_bm.sft.sft_utils.format_clean_dataset(dataset, *, mapping_func, model_family_name, base_model_hf, num_workers_format_dataset=32, process_img=False, save_processed_img_to_disk=False, new_shape_hw=None)

Format a loaded+split MedVision dataset for SFT and prune it to the training columns.

Second stage of prepare_dataset(), split out so callers can run load_split_limit_dataset() first, read the true split sizes (e.g. to name the prepared-dataset directory), and only then pay for formatting. Formatting is a row-preserving map, so the sizes do not change here.

Parameters:

• dataset (DatasetDict) – Train/validation splits as returned by load_split_limit_dataset().

• mapping_func

• model_family_name

• base_model_hf

• num_workers_format_dataset

:param : :param process_img: As in

prepare_dataset().

Parameters:

• save_processed_img_to_disk – As in prepare_dataset().

• new_shape_hw – As in prepare_dataset().

Returns:

The formatted splits containing only the retained columns.

Return type:

DatasetDict

medvision_bm.sft.sft_utils.prepare_dataset(*, tasks_list_json_path, limit_train_sample, limit_val_sample, mapping_func, model_family_name, base_model_hf, num_workers_concat_datasets=4, num_workers_format_dataset=32, tag_ds=None, process_img=False, save_processed_img_to_disk=False, new_shape_hw=None, download_mode='reuse_dataset_if_exists')

Load, format, and prune a MedVision dataset for SFT in one call.

Combines load_split_limit_dataset(), format_dataset(), and clean_dataset(): it loads and splits the tasks, maps each example into the chat messages format via mapping_func, then keeps only the columns needed for training.

Parameters:

• tasks_list_json_path (str) – Path to the JSON file listing the tasks.

• limit_train_sample (int) – Training-sample cap (< 0 = no limit, > 0 = cap).

• limit_val_sample (int) – Target validation-split size; must be > 0.

• mapping_func – Per-example formatting function applied during mapping.

• model_family_name (str) – Model family name passed to mapping_func as model_name (used for image-resize logic).

• base_model_hf (str) – HuggingFace model id passed to mapping_func as model_hf.

• num_workers_concat_datasets (int) – Worker processes for loading. Defaults to 4.

• num_workers_format_dataset (int) – Worker processes for formatting. Defaults to 32.

• tag_ds (str) – Tag embedded in task names; required by the loader.

• process_img (bool) – If True, embed processed PNG images in the dataset (processed_images column). Not recommended (large cache).

• save_processed_img_to_disk (bool) – If True, write processed PNGs to disk and store their paths in image_file_png.

• new_shape_hw (tuple[int, int] | None) – Optional explicit (height, width) to resize slices to before formatting.

• download_mode (str) – Passed through to the loader. Defaults to "reuse_dataset_if_exists".

Returns:

Train/validation splits containing only the retained columns (messages, labels, image_file, slice_dim, slice_idx, plus processed_images and/or image_file_png when enabled).

Return type:

DatasetDict

medvision_bm.sft.sft_utils.recompute_total_max_steps(trainer)

Recompute total planned update steps based on global dataset size, world size and desired epochs.

medvision_bm.sft.sft_utils.prepare_trainer(*, run_name, base_model_hf, lora_checkpoint_dir, data, make_collate_fn, per_device_train_batch_size=14, per_device_eval_batch_size=14, gradient_accumulation_steps=6, use_flash_attention_2=True, num_train_epochs=1, save_steps=100, eval_steps=50, logging_steps=50, save_total_limit=10, dataloader_num_workers=8, gradient_checkpointing=False, dataloader_pin_memory=True, push_LoRA=False, enable_temperature_sampler=False, temperature_sampler_T=3.0, temperature_sampler_task_column='__task_name', temperature_sampler_num_samples=-1, archive_every_n_steps=0)

Build a QLoRA SFTTrainer for MedVision SFT.

Loads base_model_hf in 4-bit NF4 quantization, attaches a LoRA adapter (all-linear target modules, with lm_head and embed_tokens also trained), and wraps it in an SFTTrainer configured for BF16 training and Weights & Biases logging. When enable_temperature_sampler is set, a temperature-weighted multi-task sampler is used instead of uniform sampling.

Parameters:

• run_name (str) – Run name for logging / W&B.

• base_model_hf (str) – HuggingFace id of the base image-text-to-text model.

• lora_checkpoint_dir (str) – Output directory for adapter checkpoints.

• data – DatasetDict with "train" and "validation" splits.

• make_collate_fn – Factory called as make_collate_fn(processor) to build the data collator.

• per_device_train_batch_size (int) – Per-device train batch size. Defaults to 14.

• per_device_eval_batch_size (int) – Per-device eval batch size. Defaults to 14.

• gradient_accumulation_steps (int) – Gradient accumulation steps. Defaults to 6.

• use_flash_attention_2 (bool) – Use FlashAttention-2 when True, else eager attention. Defaults to True.

• num_train_epochs (int) – Number of training epochs. Defaults to 1.

• save_steps (int) – Steps between checkpoint saves. Defaults to 100.

• eval_steps (int) – Steps between evaluations. Defaults to 50.

• logging_steps (int) – Steps between log entries. Defaults to 50.

• save_total_limit (int) – Max checkpoints to retain. Defaults to 10.

• dataloader_num_workers (int) – DataLoader worker processes. Defaults to 8.

• gradient_checkpointing (bool) – Enable gradient checkpointing. Defaults to False.

• dataloader_pin_memory (bool) – Pin DataLoader memory. Defaults to True.

• push_LoRA (bool) – Push the adapter to the Hub (private) when True. Defaults to False.

• enable_temperature_sampler (bool) – Enable temperature-based multi-task sampling. Defaults to False.

• temperature_sampler_T (float) – Sampling temperature. Defaults to 3.0.

• temperature_sampler_task_column (str) – Column holding task labels. Defaults to "__task_name".

• temperature_sampler_num_samples (int) – Draws per epoch; <= 0 keeps the training-set length. Defaults to -1.

• archive_every_n_steps (int) – When > 0, save a checkpoint every N steps and keep it under <output_dir>/archive (exempt from save_total_limit). Defaults to 0 (off).

Returns:

The configured trainer (a temperature-sampling subclass when enabled with more than one task).

Return type:

SFTTrainer

Raises:

ValueError – If the GPU does not support bfloat16.

medvision_bm.sft.sft_utils.prepare_trainer_fullFT(*, run_name, base_model_hf, model_weights_from=None, checkpoint_dir, data, make_collate_fn, per_device_train_batch_size=1, per_device_eval_batch_size=1, gradient_accumulation_steps=16, use_flash_attention_2=True, num_train_epochs=1, save_steps=100, eval_steps=50, logging_steps=50, save_total_limit=5, dataloader_num_workers=4, gradient_checkpointing=True, dataloader_pin_memory=True, push_model=False, enable_temperature_sampler=False, temperature_sampler_T=3.0, temperature_sampler_task_column='__task_name', temperature_sampler_num_samples=-1, archive_every_n_steps=0)

Prepare an SFTTrainer for full parameter finetuning (no LoRA, no quantization).

Loads the model in BF16 without any PEFT adapter. All parameters are trained. Use a lower learning rate and cosine scheduler compared to the LoRA variant. When enable_temperature_sampler is set, a temperature-weighted multi-task sampler is used instead of uniform sampling.

Parameters:

• run_name (str) – Run name for logging / W&B.

• base_model_hf (str) – HuggingFace id of the base image-text-to-text model. Also the source of the processor.

• model_weights_from (str, optional) – Local checkpoint directory to load the model weights from instead of base_model_hf (resume path). Pass the same directory to train_resume_from_checkpoint with weights_preloaded=True so the Trainer does not re-load it.

• checkpoint_dir (str) – Output directory for checkpoints.

• data – DatasetDict with "train" and "validation" splits.

• make_collate_fn – Factory called as make_collate_fn(processor) to build the data collator.

• per_device_train_batch_size (int) – Per-device train batch size. Defaults to 1.

• per_device_eval_batch_size (int) – Per-device eval batch size. Defaults to 1.

• gradient_accumulation_steps (int) – Gradient accumulation steps. Defaults to 16.

• use_flash_attention_2 (bool) – Use FlashAttention-2 when True, else eager attention. Defaults to True.

• num_train_epochs (int) – Number of training epochs. Defaults to 1.

• save_steps (int) – Steps between checkpoint saves. Defaults to 100.

• eval_steps (int) – Steps between evaluations. Defaults to 50.

• logging_steps (int) – Steps between log entries. Defaults to 50.

• save_total_limit (int) – Max checkpoints to retain. Defaults to 5.

• dataloader_num_workers (int) – DataLoader worker processes. Defaults to 4.

• gradient_checkpointing (bool) – Enable gradient checkpointing (on by default; required at 7B+ scale). Defaults to True.

• dataloader_pin_memory (bool) – Pin DataLoader memory. Defaults to True.

• push_model (bool) – Push the trained model to the Hub (private) when True. Defaults to False.

• enable_temperature_sampler (bool) – Enable temperature-based multi-task sampling. Defaults to False.

• temperature_sampler_T (float) – Sampling temperature. Defaults to 3.0.

• temperature_sampler_task_column (str) – Column holding task labels. Defaults to "__task_name".

• temperature_sampler_num_samples (int) – Draws per epoch; <= 0 keeps the training-set length. Defaults to -1.

• archive_every_n_steps (int) – When > 0, save a checkpoint every N steps and keep it under <output_dir>/archive (exempt from save_total_limit). Defaults to 0 (off).

Returns:

The configured trainer (a temperature-sampling subclass when enabled with more than one task).

Return type:

SFTTrainer

Raises:

ValueError – If the GPU does not support bfloat16.

medvision_bm.sft.sft_utils.merge_models(base_model_hf, lora_checkpoint_dir, merged_model_hf, merged_model_dir, push_to_hub)

Merge a LoRA adapter into its base model and optionally save/push it.

Loads the base model on CPU in fp32 (so the sub-BF16 LoRA delta is representable), merges the adapter with safe_merge=True, then optionally saves the merged model locally and/or pushes it to the Hugging Face Hub. This function is intended to be called only on the main process.

Parameters:

• base_model_hf (str) – HuggingFace id of the base model.

• lora_checkpoint_dir (str) – Directory of the trained LoRA adapter (also the source of the processor).

• merged_model_hf (str) – Target Hub repo id for the merged model. Required only when push_to_hub is True.

• merged_model_dir (str | None) – Local directory to save the merged model to; skipped when None.

• push_to_hub (bool) – If True, push the merged model and processor to the Hub as a private repo.

Raises:

ValueError – If push_to_hub is True but merged_model_hf is None.

medvision_bm.sft.sft_utils.fsdp_flat_param_dtype(model)

Return the dtype of the model’s FSDP flat parameters, or None.

None means the model is not FSDP-wrapped (nothing to reconcile) or its flat params have mixed dtypes (ambiguous; a warning is printed and the caller should leave the optimizer state as loaded).

medvision_bm.sft.sft_utils.recast_optimizer_state(optimizer, target_dtype)

Recast the optimizer’s floating-point moment tensors to target_dtype.

Every state tensor except step (exp_avg, exp_avg_sq, …) is moved to target_dtype on its parameter’s device; step is left exactly as loaded (fused AdamW keeps it fp32). Accepts a raw torch optimizer or accelerate’s AcceleratedOptimizer wrapper.

Why this exists: on an FSDP resume the Trainer loads the optimizer state BEFORE the first forward. With use_orig_params=True (the accelerate launch default) the optimizer holds the ORIGINAL parameter views created at wrap time, still in the dtype the model was loaded in (bf16); accelerate’s mixed-precision upcast only replaces the FLAT params’ data with fp32, and FSDP refreshes the views on the first forward (lazy_init). Optimizer.load_state_dict therefore casts the moments to the stale bf16 dtype and the first fused-AdamW step sees fp32 params next to bf16 moments: RuntimeError: Tensors of the same index must be on the same device and the same dtype (observed 2026-09-21 resuming Qwen2.5-VL-7B fullFT from checkpoint-1000). A fresh run never hits this because its moments are created lazily AFTER the refresh.

Returns:

Number of tensors recast.

Return type:

int

medvision_bm.sft.sft_utils.match_checkpoint_global_batch(trainer, last_checkpoint)

Rescale gradient_accumulation_steps so the global batch equals the checkpoint’s.

HF Trainer counts optimizer steps. Resuming on a different number of GPUs with the same gradient_accumulation_steps changes the samples per step, but the restored global_step is not rescaled: the epoch counter, LR schedule position and resume data-skip all jump (2 -> 4 GPUs moved a run from epoch 0.41 to 0.83 and skipped the data in between). Keeping per_device_bsz * world_size * grad_accum equal to the checkpoint’s global batch keeps every step-indexed quantity valid.

The checkpoint’s global batch is inferred from trainer_state.json as epoch * len(train_dataset) / global_step (the same samples-per-epoch definition recompute_total_max_steps uses), rounded to a multiple of its per-device batch.

medvision_bm.sft.sft_utils.train_resume_from_checkpoint(trainer, last_checkpoint, weights_preloaded=False)

medvision_bm.sft.sft_utils.parse_args_multiTask()

Parse command-line arguments for SFT on the MedVision dataset.

medvision_bm.sft.sft_utils.check_model_supported(model_name)

medvision_bm.sft.sft_utils.parse_validate_args_multiTask()

medvision_bm.sft.sft_utils.parse_sample_limits(**kwargs)

Determine sample limits for each task with fallbacks.

Logic:

• Missing/None limit kwargs fall back to -1 (train: no limit, use the full dataset) or 100 (val fallbacks), mirroring the argparse defaults

• If task-specific limit > 0: use it

• Else: use per-task limit

• An explicit limit of 0 raises ValueError (ambiguous): to skip a task, omit its tasks_list_json_path instead

• If task JSON path is None: set limit to 0 (task not used)

Returns:

(train_limit_AD, val_limit_AD,

train_limit_detect, val_limit_detect, train_limit_TL, val_limit_TL, train_limit_total)

Return type:

A tuple of sample limits

medvision_bm.sft.sft_utils.limit_prepared_validation_split(dataset, *, per_task_limits, total_limit, task_column, seed=1024)

Apply the validation sample limits to an already-PREPARED validation split.

Dataset preparation carves each task’s validation split with the val_sample_limit* knobs, but a run that loads a prepared dataset (--skip_process_dataset with --prepared_ds_dir) bypasses preparation and would otherwise evaluate on the split exactly as saved (1.8M rows for the full v1.4.0 data). This re-applies the same knobs to the loaded split so periodic evaluation stays affordable:

1. per task: rows whose task_column equals a key of per_task_limits with a positive limit are sub-sampled down to that limit (uniformly, without replacement); a task at or below its limit, or without a positive limit, is kept whole, and the surviving rows keep the split’s original order;

2. total: if total_limit is positive and smaller than what remains, the split is shuffled with seed and truncated to total_limit.

Rows are only ever dropped, never upsampled. Every selection is seeded so all ranks keep identical rows. If task_column is missing, the per-task step is skipped with a warning and only the total limit applies. A split already within its limits (e.g. one prepared by the same launcher) is returned unchanged.

Parameters:

• dataset (datasets.Dataset) – The loaded validation split.

• per_task_limits (dict[str, int | None]) – Task label -> row limit (<= 0 / None keeps that task whole).

• total_limit (int | None) – Overall row limit (<= 0 / None keeps all).

• task_column (str) – Column holding each row’s task label.

• seed (int) – RNG seed. Defaults to the project SEED.

Returns:

The (possibly) reduced validation split.

Return type:

datasets.Dataset

medvision_bm.sft.sft_utils.mask_non_assistant_turns(input_ids, labels, tokenizer)

Mask everything except assistant response content + its closing <|im_end|>.

Completion-only masking: for every assistant turn the header <|im_start|>assistant\n is masked (it is chat-template scaffolding the model never needs to generate), along with all system/user/tool turns; loss is computed only on the response tokens and the <|im_end|> that terminates the assistant turn. The trailing newline after <|im_end|> is also masked.

ChatML only (Qwen). Gemma checkpoints must use mask_non_assistant_turns_gemma.

medvision_bm.sft.sft_utils.mask_non_assistant_turns_gemma(input_ids, labels, tokenizer)

Gemma-family counterpart of mask_non_assistant_turns().

Same completion-only policy, but the turn markers come from the Gemma vocabulary. Gemma 3 / MedGemma render <start_of_turn>model\n ... <end_of_turn>; Gemma 4 renders <|turn>model\n ... <turn|>. Both use the single token “model” as the assistant role.

The families must not share a masker: on a Gemma tokenizer <|im_start|> resolves to unk_token_id, no assistant header is ever matched, every label becomes -100 and the loss silently becomes NaN. Hence the probe below refuses rather than guesses.