⚡️ Speed up function create_finetune_request by 15%

[codeflash-ai]

📄 15% (0.15x) speedup for create_finetune_request in src/together/resources/finetune.py

⏱️ Runtime : 2.26 milliseconds → 1.97 milliseconds (best of 88 runs)

📝 Explanation and details

The optimized code achieves a 15% speedup through two key optimizations:

1. Fast-path optimization in log_warn_once:

• Added an early membership check using a lightweight string key (msg_candidate) before expensive logfmt formatting
• This avoids costly regex operations and dictionary formatting for duplicate warnings (the common case)
• From profiler: reduces logfmt calls from 44 to 1, saving ~550μs per duplicate warning

2. Reduced attribute lookups in create_finetune_request:

• Cached model_limits.lora_training and model_limits.full_training in local variables (lora_cfg, full_cfg)
• Eliminated repeated attribute access overhead when extracting batch size limits
• Consolidated validation logic to use cached references instead of multiple dotted lookups

Performance characteristics by test case:

• Basic cases (20-38% faster): Benefit most from reduced attribute lookups during normal execution
• Error cases with DPO/SFT validation (90-200% faster): Fast-path through validation logic with fewer attribute accesses
• Large scale cases (20-35% faster): Compound benefits from both optimizations
• Warning-heavy cases (minimal impact): The log_warn_once optimization only helps on repeated calls

The optimizations maintain identical behavior while eliminating redundant work - particularly effective for the common path of successful request creation and duplicate warning suppression.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	✅ 32 Passed
🌀 Generated Regression Tests	✅ 79 Passed
⏪ Replay Tests	✅ 4 Passed
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	98.8%

⚙️ Existing Unit Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
unit/test_finetune_resources.py::test_bad_max_grad_norm	12.9μs	12.6μs	2.22%✅
unit/test_finetune_resources.py::test_bad_min_lr_ratio	11.5μs	11.3μs	2.45%✅
unit/test_finetune_resources.py::test_bad_training_method	6.46μs	8.26μs	-21.8%⚠️
unit/test_finetune_resources.py::test_bad_warmup	11.8μs	12.2μs	-3.23%⚠️
unit/test_finetune_resources.py::test_bad_weight_decay	12.9μs	13.2μs	-2.74%⚠️
unit/test_finetune_resources.py::test_batch_size_limit	34.5μs	29.6μs	16.3%✅
unit/test_finetune_resources.py::test_both_from_checkpoint_model_name	1.12μs	1.13μs	-0.976%⚠️
unit/test_finetune_resources.py::test_dpo_request	22.9μs	23.7μs	-3.25%⚠️
unit/test_finetune_resources.py::test_dpo_request_lora	29.6μs	28.1μs	5.06%✅
unit/test_finetune_resources.py::test_from_checkpoint_request	30.8μs	24.3μs	26.5%✅
unit/test_finetune_resources.py::test_lora_request	39.0μs	29.0μs	34.6%✅
unit/test_finetune_resources.py::test_lora_request_with_lora_dropout	70.9μs	54.8μs	29.4%✅
unit/test_finetune_resources.py::test_no_from_checkpoint_no_model_name	1.13μs	1.12μs	0.806%✅
unit/test_finetune_resources.py::test_no_training_file	2.68μs	2.70μs	-0.777%⚠️
unit/test_finetune_resources.py::test_non_full_model	3.44μs	1.83μs	87.7%✅
unit/test_finetune_resources.py::test_non_lora_model	3.12μs	1.70μs	83.3%✅
unit/test_finetune_resources.py::test_simple_request	26.7μs	21.2μs	25.8%✅
unit/test_finetune_resources.py::test_train_on_inputs_for_sft	87.2μs	80.8μs	7.84%✅
unit/test_finetune_resources.py::test_train_on_inputs_not_supported_for_dpo	5.28μs	5.48μs	-3.74%⚠️
unit/test_finetune_resources.py::test_validation_file	33.6μs	27.6μs	21.9%✅

🌀 Generated Regression Tests and Runtime

import pytest
from together.resources.finetune import create_finetune_request

# --- Minimal stubs for required types/classes ---

class CosineLRSchedulerArgs:
    def __init__(self, min_lr_ratio, num_cycles):
        self.min_lr_ratio = min_lr_ratio
        self.num_cycles = num_cycles

class CosineLRScheduler:
    def __init__(self, lr_scheduler_args):
        self.lr_scheduler_args = lr_scheduler_args

class LinearLRSchedulerArgs:
    def __init__(self, min_lr_ratio):
        self.min_lr_ratio = min_lr_ratio

class LinearLRScheduler:
    def __init__(self, lr_scheduler_args):
        self.lr_scheduler_args = lr_scheduler_args

class FullTrainingType:
    pass

class LoRATrainingType:
    def __init__(self, lora_r, lora_alpha, lora_dropout, lora_trainable_modules):
        self.lora_r = lora_r
        self.lora_alpha = lora_alpha
        self.lora_dropout = lora_dropout
        self.lora_trainable_modules = lora_trainable_modules

class TrainingMethodSFT:
    def __init__(self, train_on_inputs=None):
        self.method = "sft"
        self.train_on_inputs = train_on_inputs

class TrainingMethodDPO:
    def __init__(self, dpo_beta=None, dpo_normalize_logratios_by_length=False, dpo_reference_free=False, rpo_alpha=None, simpo_gamma=None):
        self.method = "dpo"
        self.dpo_beta = dpo_beta
        self.dpo_normalize_logratios_by_length = dpo_normalize_logratios_by_length
        self.dpo_reference_free = dpo_reference_free
        self.rpo_alpha = rpo_alpha
        self.simpo_gamma = simpo_gamma

class FinetuneRequest:
    def __init__(self, **kwargs):
        self.__dict__.update(kwargs)

class FinetuneTrainingLimits:
    def __init__(self, full_training=None, lora_training=None):
        self.full_training = full_training
        self.lora_training = lora_training

class TrainingLimit:
    def __init__(self, max_batch_size, min_batch_size, max_batch_size_dpo, max_rank=None):
        self.max_batch_size = max_batch_size
        self.min_batch_size = min_batch_size
        self.max_batch_size_dpo = max_batch_size_dpo
        self.max_rank = max_rank
from together.resources.finetune import create_finetune_request

# --- Unit tests ---

# Helper: default training limits for basic tests
default_full_limits = TrainingLimit(max_batch_size=64, min_batch_size=1, max_batch_size_dpo=32)
default_lora_limits = TrainingLimit(max_batch_size=32, min_batch_size=1, max_batch_size_dpo=16, max_rank=8)

# 1. Basic Test Cases

def test_basic_sft_full_training():
    """Basic SFT full training request with all required arguments."""
    limits = FinetuneTrainingLimits(full_training=default_full_limits)
    codeflash_output = create_finetune_request(
        model_limits=limits,
        training_file="train.jsonl",
        model="base-model"
    ); req = codeflash_output # 33.7μs -> 27.7μs (21.7% faster)

def test_basic_lora_training():
    """Basic LoRA training request with LoRA args."""
    limits = FinetuneTrainingLimits(full_training=default_full_limits, lora_training=default_lora_limits)
    codeflash_output = create_finetune_request(
        model_limits=limits,
        training_file="train.jsonl",
        model="base-model",
        lora=True,
        lora_r=4,
        lora_alpha=8,
        lora_dropout=0.5,
        lora_trainable_modules="linear-only"
    ); req = codeflash_output # 32.3μs -> 25.0μs (29.5% faster)

def test_basic_dpo_training():
    """Basic DPO training request with DPO args."""
    limits = FinetuneTrainingLimits(full_training=default_full_limits)
    codeflash_output = create_finetune_request(
        model_limits=limits,
        training_file="train.jsonl",
        model="base-model",
        training_method="dpo",
        dpo_beta=0.1
    ); req = codeflash_output # 22.7μs -> 22.6μs (0.336% faster)

def test_basic_batch_size_within_limits():
    """Batch size within allowed limits."""
    limits = FinetuneTrainingLimits(full_training=default_full_limits)
    codeflash_output = create_finetune_request(
        model_limits=limits,
        training_file="train.jsonl",
        model="base-model",
        batch_size=10
    ); req = codeflash_output # 26.6μs -> 20.5μs (29.2% faster)

def test_basic_cosine_scheduler():
    """Cosine scheduler with valid cycles."""
    limits = FinetuneTrainingLimits(full_training=default_full_limits)
    codeflash_output = create_finetune_request(
        model_limits=limits,
        training_file="train.jsonl",
        model="base-model",
        lr_scheduler_type="cosine",
        scheduler_num_cycles=1.0
    ); req = codeflash_output # 25.8μs -> 19.9μs (29.6% faster)

def test_basic_linear_scheduler():
    """Linear scheduler selected."""
    limits = FinetuneTrainingLimits(full_training=default_full_limits)
    codeflash_output = create_finetune_request(
        model_limits=limits,
        training_file="train.jsonl",
        model="base-model",
        lr_scheduler_type="linear"
    ); req = codeflash_output # 23.0μs -> 17.8μs (29.3% faster)

# 2. Edge Test Cases

def test_error_both_model_and_checkpoint():
    """Error if both model and checkpoint are specified."""
    limits = FinetuneTrainingLimits(full_training=default_full_limits)
    with pytest.raises(ValueError, match="either a model or a checkpoint to start a job from, not both"):
        create_finetune_request(
            model_limits=limits,
            training_file="train.jsonl",
            model="base-model",
            from_checkpoint="checkpoint"
        ) # 1.19μs -> 1.23μs (3.34% slower)

def test_error_neither_model_nor_checkpoint():
    """Error if neither model nor checkpoint is specified."""
    limits = FinetuneTrainingLimits(full_training=default_full_limits)
    with pytest.raises(ValueError, match="You must specify either a model or a checkpoint"):
        create_finetune_request(
            model_limits=limits,
            training_file="train.jsonl"
        ) # 1.13μs -> 1.17μs (3.43% slower)



def test_error_lora_not_supported():
    """Error if LoRA requested but not supported."""
    limits = FinetuneTrainingLimits(full_training=default_full_limits)
    with pytest.raises(ValueError, match="LoRA adapters are not supported"):
        create_finetune_request(
            model_limits=limits,
            training_file="train.jsonl",
            model="base-model",
            lora=True
        ) # 8.03μs -> 2.02μs (298% faster)

def test_error_full_training_not_supported():
    """Error if full training requested but not supported."""
    limits = FinetuneTrainingLimits(lora_training=default_lora_limits)
    with pytest.raises(ValueError, match="Full training is not supported"):
        create_finetune_request(
            model_limits=limits,
            training_file="train.jsonl",
            model="base-model"
        ) # 4.84μs -> 1.67μs (189% faster)

def test_error_batch_size_too_high_sft():
    """Error if batch size is above max for SFT."""
    limits = FinetuneTrainingLimits(full_training=default_full_limits)
    with pytest.raises(ValueError, match="Requested batch size .* higher that the maximum allowed value"):
        create_finetune_request(
            model_limits=limits,
            training_file="train.jsonl",
            model="base-model",
            batch_size=100
        ) # 5.35μs -> 8.38μs (36.2% slower)

def test_error_batch_size_too_high_dpo():
    """Error if batch size is above max for DPO."""
    limits = FinetuneTrainingLimits(full_training=default_full_limits)
    with pytest.raises(ValueError, match="Requested batch size .* higher that the maximum allowed value"):
        create_finetune_request(
            model_limits=limits,
            training_file="train.jsonl",
            model="base-model",
            training_method="dpo",
            batch_size=100
        ) # 5.40μs -> 5.70μs (5.30% slower)

def test_error_batch_size_too_low():
    """Error if batch size is below min."""
    limits = FinetuneTrainingLimits(full_training=default_full_limits)
    with pytest.raises(ValueError, match="Requested batch size .* lower that the minimum allowed value"):
        create_finetune_request(
            model_limits=limits,
            training_file="train.jsonl",
            model="base-model",
            batch_size=0
        ) # 5.03μs -> 4.99μs (0.782% faster)

def test_error_warmup_ratio_out_of_bounds():
    """Error if warmup_ratio is out of [0,1]."""
    limits = FinetuneTrainingLimits(full_training=default_full_limits)
    with pytest.raises(ValueError, match="Warmup ratio should be between 0 and 1"):
        create_finetune_request(
            model_limits=limits,
            training_file="train.jsonl",
            model="base-model",
            warmup_ratio=2.0
        ) # 5.71μs -> 5.86μs (2.63% slower)

def test_error_min_lr_ratio_out_of_bounds():
    """Error if min_lr_ratio is out of [0,1]."""
    limits = FinetuneTrainingLimits(full_training=default_full_limits)
    with pytest.raises(ValueError, match="Min learning rate ratio should be between 0 and 1"):
        create_finetune_request(
            model_limits=limits,
            training_file="train.jsonl",
            model="base-model",
            min_lr_ratio=-0.1
        ) # 6.67μs -> 6.88μs (3.05% slower)

def test_error_max_grad_norm_negative():
    """Error if max_grad_norm is negative."""
    limits = FinetuneTrainingLimits(full_training=default_full_limits)
    with pytest.raises(ValueError, match="Max gradient norm should be non-negative"):
        create_finetune_request(
            model_limits=limits,
            training_file="train.jsonl",
            model="base-model",
            max_grad_norm=-1.0
        ) # 5.67μs -> 5.93μs (4.55% slower)

def test_error_weight_decay_negative():
    """Error if weight_decay is negative."""
    limits = FinetuneTrainingLimits(full_training=default_full_limits)
    with pytest.raises(ValueError, match="Weight decay should be non-negative"):
        create_finetune_request(
            model_limits=limits,
            training_file="train.jsonl",
            model="base-model",
            weight_decay=-0.1
        ) # 7.03μs -> 7.11μs (1.13% slower)

def test_error_invalid_training_method():
    """Error if training_method is not recognized."""
    limits = FinetuneTrainingLimits(full_training=default_full_limits)
    with pytest.raises(ValueError, match="training_method must be one of"):
        create_finetune_request(
            model_limits=limits,
            training_file="train.jsonl",
            model="base-model",
            training_method="foobar"
        ) # 6.28μs -> 6.37μs (1.40% slower)

def test_error_train_on_inputs_for_dpo():
    """Error if train_on_inputs is set for DPO."""
    limits = FinetuneTrainingLimits(full_training=default_full_limits)
    with pytest.raises(ValueError, match="train_on_inputs is only supported for SFT training"):
        create_finetune_request(
            model_limits=limits,
            training_file="train.jsonl",
            model="base-model",
            training_method="dpo",
            train_on_inputs=True
        ) # 5.30μs -> 5.61μs (5.54% slower)

def test_error_dpo_beta_for_sft():
    """Error if dpo_beta is set for SFT."""
    limits = FinetuneTrainingLimits(full_training=default_full_limits)
    with pytest.raises(ValueError, match="dpo_beta is only supported for DPO training"):
        create_finetune_request(
            model_limits=limits,
            training_file="train.jsonl",
            model="base-model",
            dpo_beta=0.1
        ) # 13.8μs -> 6.51μs (112% faster)

def test_error_dpo_normalize_logratios_for_sft():
    """Error if dpo_normalize_logratios_by_length is set for SFT."""
    limits = FinetuneTrainingLimits(full_training=default_full_limits)
    with pytest.raises(ValueError, match="dpo_normalize_logratios_by_length=True is only supported for DPO training"):
        create_finetune_request(
            model_limits=limits,
            training_file="train.jsonl",
            model="base-model",
            dpo_normalize_logratios_by_length=True
        ) # 11.9μs -> 5.97μs (99.1% faster)

def test_error_rpo_alpha_for_sft():
    """Error if rpo_alpha is set for SFT."""
    limits = FinetuneTrainingLimits(full_training=default_full_limits)
    with pytest.raises(ValueError, match="rpo_alpha is only supported for DPO training"):
        create_finetune_request(
            model_limits=limits,
            training_file="train.jsonl",
            model="base-model",
            rpo_alpha=0.5
        ) # 11.7μs -> 6.06μs (92.9% faster)

def test_error_rpo_alpha_negative():
    """Error if rpo_alpha is negative for DPO."""
    limits = FinetuneTrainingLimits(full_training=default_full_limits)
    with pytest.raises(ValueError, match="rpo_alpha should be non-negative"):
        create_finetune_request(
            model_limits=limits,
            training_file="train.jsonl",
            model="base-model",
            training_method="dpo",
            rpo_alpha=-0.1
        ) # 6.72μs -> 7.11μs (5.49% slower)

def test_error_simpo_gamma_for_sft():
    """Error if simpo_gamma is set for SFT."""
    limits = FinetuneTrainingLimits(full_training=default_full_limits)
    with pytest.raises(ValueError, match="simpo_gamma is only supported for DPO training"):
        create_finetune_request(
            model_limits=limits,
            training_file="train.jsonl",
            model="base-model",
            simpo_gamma=0.5
        ) # 12.1μs -> 6.24μs (93.7% faster)

def test_error_simpo_gamma_negative():
    """Error if simpo_gamma is negative for DPO."""
    limits = FinetuneTrainingLimits(full_training=default_full_limits)
    with pytest.raises(ValueError, match="simpo_gamma should be non-negative"):
        create_finetune_request(
            model_limits=limits,
            training_file="train.jsonl",
            model="base-model",
            training_method="dpo",
            simpo_gamma=-0.1
        ) # 6.64μs -> 6.68μs (0.614% slower)

def test_error_scheduler_num_cycles_zero():
    """Error if scheduler_num_cycles is zero for cosine."""
    limits = FinetuneTrainingLimits(full_training=default_full_limits)
    with pytest.raises(ValueError, match="Number of cycles should be greater than 0"):
        create_finetune_request(
            model_limits=limits,
            training_file="train.jsonl",
            model="base-model",
            lr_scheduler_type="cosine",
            scheduler_num_cycles=0.0
        ) # 13.1μs -> 7.28μs (80.2% faster)

def test_error_lora_dropout_out_of_range():
    """Error if lora_dropout is not in [0,1)."""
    limits = FinetuneTrainingLimits(full_training=default_full_limits, lora_training=default_lora_limits)
    with pytest.raises(ValueError, match="LoRA dropout must be in \[0, 1\) range"):
        create_finetune_request(
            model_limits=limits,
            training_file="train.jsonl",
            model="base-model",
            lora=True,
            lora_dropout=1.0
        ) # 4.87μs -> 1.62μs (200% faster)

# 3. Large Scale Test Cases

def test_large_batch_size_at_max():
    """Test with batch size at maximum allowed."""
    limits = FinetuneTrainingLimits(full_training=TrainingLimit(max_batch_size=999, min_batch_size=1, max_batch_size_dpo=500))
    codeflash_output = create_finetune_request(
        model_limits=limits,
        training_file="train.jsonl",
        model="base-model",
        batch_size=999
    ); req = codeflash_output # 35.9μs -> 30.6μs (17.4% faster)

def test_large_batch_size_at_max_dpo():
    """Test with batch size at maximum allowed for DPO."""
    limits = FinetuneTrainingLimits(full_training=TrainingLimit(max_batch_size=999, min_batch_size=1, max_batch_size_dpo=500))
    codeflash_output = create_finetune_request(
        model_limits=limits,
        training_file="train.jsonl",
        model="base-model",
        training_method="dpo",
        batch_size=500
    ); req = codeflash_output # 23.5μs -> 25.3μs (7.13% slower)

def test_large_lora_rank_and_alpha():
    """Test with large lora_r and lora_alpha values."""
    limits = FinetuneTrainingLimits(full_training=default_full_limits, lora_training=TrainingLimit(max_batch_size=32, min_batch_size=1, max_batch_size_dpo=16, max_rank=999))
    codeflash_output = create_finetune_request(
        model_limits=limits,
        training_file="train.jsonl",
        model="base-model",
        lora=True,
        lora_r=999,
        lora_alpha=1998
    ); req = codeflash_output # 33.7μs -> 25.5μs (32.3% faster)

def test_large_num_epochs():
    """Test with large n_epochs value."""
    limits = FinetuneTrainingLimits(full_training=default_full_limits)
    codeflash_output = create_finetune_request(
        model_limits=limits,
        training_file="train.jsonl",
        model="base-model",
        n_epochs=999
    ); req = codeflash_output # 26.2μs -> 20.5μs (27.6% faster)

def test_large_scheduler_num_cycles():
    """Test with large scheduler_num_cycles for cosine."""
    limits = FinetuneTrainingLimits(full_training=default_full_limits)
    codeflash_output = create_finetune_request(
        model_limits=limits,
        training_file="train.jsonl",
        model="base-model",
        lr_scheduler_type="cosine",
        scheduler_num_cycles=999.0
    ); req = codeflash_output # 25.4μs -> 20.2μs (25.7% faster)

def test_large_weight_decay():
    """Test with large weight_decay value."""
    limits = FinetuneTrainingLimits(full_training=default_full_limits)
    codeflash_output = create_finetune_request(
        model_limits=limits,
        training_file="train.jsonl",
        model="base-model",
        weight_decay=999.0
    ); req = codeflash_output # 25.7μs -> 19.6μs (31.5% faster)

def test_large_scale_dpo_with_simpo_gamma():
    """Test large scale DPO with simpo_gamma > 0 triggers reference-free and normalization."""
    limits = FinetuneTrainingLimits(full_training=default_full_limits)
    codeflash_output = create_finetune_request(
        model_limits=limits,
        training_file="train.jsonl",
        model="base-model",
        training_method="dpo",
        simpo_gamma=0.5
    ); req = codeflash_output # 326μs -> 324μs (0.779% faster)

def test_large_scale_all_fields():
    """Test with all optional fields set to large values."""
    limits = FinetuneTrainingLimits(full_training=TrainingLimit(max_batch_size=999, min_batch_size=1, max_batch_size_dpo=999))
    codeflash_output = create_finetune_request(
        model_limits=limits,
        training_file="train.jsonl",
        model="base-model",
        validation_file="val.jsonl",
        n_epochs=999,
        n_evals=999,
        n_checkpoints=999,
        batch_size=999,
        learning_rate=0.999,
        lr_scheduler_type="linear",
        min_lr_ratio=1.0,
        warmup_ratio=1.0,
        max_grad_norm=999.0,
        weight_decay=999.0,
        suffix="big-job",
        wandb_api_key="key",
        wandb_base_url="url",
        wandb_project_name="proj",
        wandb_name="run",
        train_on_inputs=True,
        training_method="sft",
        from_checkpoint=None,
        from_hf_model=None,
        hf_model_revision="rev",
        hf_api_token="token",
        hf_output_repo_name="repo"
    ); req = codeflash_output # 20.2μs -> 20.8μs (2.74% slower)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import pytest
from together.resources.finetune import create_finetune_request


# Minimal stubs for required types/classes (since we can't import the actual ones)
class CosineLRSchedulerArgs:
    def __init__(self, min_lr_ratio, num_cycles):
        self.min_lr_ratio = min_lr_ratio
        self.num_cycles = num_cycles

class CosineLRScheduler:
    def __init__(self, lr_scheduler_args):
        self.lr_scheduler_args = lr_scheduler_args

class LinearLRSchedulerArgs:
    def __init__(self, min_lr_ratio):
        self.min_lr_ratio = min_lr_ratio

class LinearLRScheduler:
    def __init__(self, lr_scheduler_args):
        self.lr_scheduler_args = lr_scheduler_args

class FullTrainingType:
    pass

class LoRATrainingType:
    def __init__(self, lora_r, lora_alpha, lora_dropout, lora_trainable_modules):
        self.lora_r = lora_r
        self.lora_alpha = lora_alpha
        self.lora_dropout = lora_dropout
        self.lora_trainable_modules = lora_trainable_modules

class TrainingMethodSFT:
    def __init__(self, train_on_inputs=None):
        self.method = "sft"
        self.train_on_inputs = train_on_inputs

class TrainingMethodDPO:
    def __init__(self, dpo_beta=None, dpo_normalize_logratios_by_length=False, dpo_reference_free=False, rpo_alpha=None, simpo_gamma=None):
        self.method = "dpo"
        self.dpo_beta = dpo_beta
        self.dpo_normalize_logratios_by_length = dpo_normalize_logratios_by_length
        self.dpo_reference_free = dpo_reference_free
        self.rpo_alpha = rpo_alpha
        self.simpo_gamma = simpo_gamma

class FinetuneRequest:
    def __init__(self, **kwargs):
        self.__dict__.update(kwargs)

class FinetuneTrainingLimits:
    def __init__(self, full_training=None, lora_training=None):
        self.full_training = full_training
        self.lora_training = lora_training

class TrainingLimits:
    def __init__(self, max_batch_size, min_batch_size, max_batch_size_dpo, max_rank=None):
        self.max_batch_size = max_batch_size
        self.min_batch_size = min_batch_size
        self.max_batch_size_dpo = max_batch_size_dpo
        self.max_rank = max_rank
from together.resources.finetune import create_finetune_request

# ---------------------
# UNIT TESTS BEGIN HERE
# ---------------------

# Fixtures for model limits
@pytest.fixture
def full_training_limits():
    # max_batch_size, min_batch_size, max_batch_size_dpo
    return TrainingLimits(max_batch_size=128, min_batch_size=4, max_batch_size_dpo=64)

@pytest.fixture
def lora_training_limits():
    # max_batch_size, min_batch_size, max_batch_size_dpo, max_rank
    return TrainingLimits(max_batch_size=64, min_batch_size=2, max_batch_size_dpo=32, max_rank=16)

@pytest.fixture
def model_limits(full_training_limits, lora_training_limits):
    return FinetuneTrainingLimits(full_training=full_training_limits, lora_training=lora_training_limits)

@pytest.fixture
def full_only_model_limits(full_training_limits):
    return FinetuneTrainingLimits(full_training=full_training_limits, lora_training=None)

@pytest.fixture
def lora_only_model_limits(lora_training_limits):
    return FinetuneTrainingLimits(full_training=None, lora_training=lora_training_limits)

# 1. BASIC TEST CASES

def test_basic_full_training(model_limits):
    # Test with only required params for full training
    codeflash_output = create_finetune_request(
        model_limits=model_limits,
        training_file="train.jsonl",
        model="base-model",
    ); req = codeflash_output # 36.9μs -> 30.8μs (20.0% faster)

def test_basic_lora_training(model_limits):
    # Test with lora=True and required params
    codeflash_output = create_finetune_request(
        model_limits=model_limits,
        training_file="train.jsonl",
        model="base-model",
        lora=True,
    ); req = codeflash_output # 35.2μs -> 26.3μs (34.0% faster)

def test_checkpoint_only(model_limits):
    # Test with from_checkpoint only
    codeflash_output = create_finetune_request(
        model_limits=model_limits,
        training_file="train.jsonl",
        from_checkpoint="ckpt-123",
    ); req = codeflash_output # 29.1μs -> 22.2μs (31.6% faster)

def test_hf_model_and_base_model(model_limits):
    # Test with from_hf_model and base model
    codeflash_output = create_finetune_request(
        model_limits=model_limits,
        training_file="train.jsonl",
        model="base-model",
        from_hf_model="hf-model",
    ); req = codeflash_output # 28.6μs -> 22.0μs (30.1% faster)

def test_dpo_training_method(model_limits):
    # Test with DPO training method
    codeflash_output = create_finetune_request(
        model_limits=model_limits,
        training_file="train.jsonl",
        model="base-model",
        training_method="dpo",
    ); req = codeflash_output # 24.0μs -> 25.3μs (5.05% slower)

def test_custom_batch_size(model_limits):
    # Test with custom batch size within limits
    codeflash_output = create_finetune_request(
        model_limits=model_limits,
        training_file="train.jsonl",
        model="base-model",
        batch_size=32,
    ); req = codeflash_output # 28.4μs -> 21.9μs (29.7% faster)

def test_custom_scheduler_linear(model_limits):
    # Test with linear scheduler
    codeflash_output = create_finetune_request(
        model_limits=model_limits,
        training_file="train.jsonl",
        model="base-model",
        lr_scheduler_type="linear",
        min_lr_ratio=0.5,
    ); req = codeflash_output # 26.5μs -> 20.3μs (30.2% faster)

def test_lora_custom_params(model_limits):
    # Test lora with custom params
    codeflash_output = create_finetune_request(
        model_limits=model_limits,
        training_file="train.jsonl",
        model="base-model",
        lora=True,
        lora_r=8,
        lora_alpha=20,
        lora_dropout=0.1,
        lora_trainable_modules="custom-modules",
    ); req = codeflash_output # 32.6μs -> 25.1μs (30.2% faster)

def test_train_on_inputs_auto(model_limits):
    # train_on_inputs should default to "auto" for SFT if not specified
    codeflash_output = create_finetune_request(
        model_limits=model_limits,
        training_file="train.jsonl",
        model="base-model",
    ); req = codeflash_output # 26.9μs -> 21.6μs (24.5% faster)

# 2. EDGE TEST CASES

def test_model_and_checkpoint_error(model_limits):
    # Should raise if both model and from_checkpoint are specified
    with pytest.raises(ValueError):
        create_finetune_request(
            model_limits=model_limits,
            training_file="train.jsonl",
            model="base-model",
            from_checkpoint="ckpt-123",
        ) # 1.31μs -> 1.36μs (3.82% slower)

def test_neither_model_nor_checkpoint_error(model_limits):
    # Should raise if neither model nor from_checkpoint are specified
    with pytest.raises(ValueError):
        create_finetune_request(
            model_limits=model_limits,
            training_file="train.jsonl",
        ) # 1.26μs -> 1.26μs (0.158% slower)

def test_checkpoint_and_hf_model_error(model_limits):
    # Should raise if both from_checkpoint and from_hf_model are specified
    with pytest.raises(ValueError):
        create_finetune_request(
            model_limits=model_limits,
            training_file="train.jsonl",
            from_checkpoint="ckpt-123",
            from_hf_model="hf-model",
        ) # 1.37μs -> 1.33μs (2.93% faster)

def test_hf_model_without_base_model_error(model_limits):
    # Should raise if from_hf_model is specified without model
    with pytest.raises(ValueError):
        create_finetune_request(
            model_limits=model_limits,
            training_file="train.jsonl",
            from_hf_model="hf-model",
        ) # 1.37μs -> 1.29μs (6.10% faster)

def test_lora_not_supported(full_only_model_limits):
    # Should raise if lora=True but lora_training is None
    with pytest.raises(ValueError):
        create_finetune_request(
            model_limits=full_only_model_limits,
            training_file="train.jsonl",
            model="base-model",
            lora=True,
        ) # 6.88μs -> 1.90μs (263% faster)

def test_full_training_not_supported(lora_only_model_limits):
    # Should raise if lora=False but full_training is None
    with pytest.raises(ValueError):
        create_finetune_request(
            model_limits=lora_only_model_limits,
            training_file="train.jsonl",
            model="base-model",
            lora=False,
        ) # 5.05μs -> 1.77μs (185% faster)

def test_batch_size_too_high(model_limits):
    # Should raise if batch_size > max_batch_size for SFT
    with pytest.raises(ValueError):
        create_finetune_request(
            model_limits=model_limits,
            training_file="train.jsonl",
            model="base-model",
            batch_size=9999,
        ) # 5.58μs -> 7.87μs (29.1% slower)

def test_batch_size_too_low(model_limits):
    # Should raise if batch_size < min_batch_size
    with pytest.raises(ValueError):
        create_finetune_request(
            model_limits=model_limits,
            training_file="train.jsonl",
            model="base-model",
            batch_size=1,
        ) # 5.36μs -> 5.49μs (2.35% slower)

def test_batch_size_too_high_dpo(model_limits):
    # Should raise if batch_size > max_batch_size_dpo for DPO
    with pytest.raises(ValueError):
        create_finetune_request(
            model_limits=model_limits,
            training_file="train.jsonl",
            model="base-model",
            batch_size=9999,
            training_method="dpo",
        ) # 5.57μs -> 5.69μs (2.20% slower)

def test_warmup_ratio_out_of_bounds(model_limits):
    # Should raise if warmup_ratio > 1 or < 0
    with pytest.raises(ValueError):
        create_finetune_request(
            model_limits=model_limits,
            training_file="train.jsonl",
            model="base-model",
            warmup_ratio=1.1,
        ) # 6.87μs -> 7.15μs (3.96% slower)
    with pytest.raises(ValueError):
        create_finetune_request(
            model_limits=model_limits,
            training_file="train.jsonl",
            model="base-model",
            warmup_ratio=-0.1,
        ) # 3.14μs -> 3.21μs (2.12% slower)

def test_min_lr_ratio_out_of_bounds(model_limits):
    # Should raise if min_lr_ratio > 1 or < 0
    with pytest.raises(ValueError):
        create_finetune_request(
            model_limits=model_limits,
            training_file="train.jsonl",
            model="base-model",
            min_lr_ratio=1.5,
        ) # 5.87μs -> 6.07μs (3.22% slower)
    with pytest.raises(ValueError):
        create_finetune_request(
            model_limits=model_limits,
            training_file="train.jsonl",
            model="base-model",
            min_lr_ratio=-0.5,
        ) # 3.15μs -> 3.23μs (2.42% slower)

def test_max_grad_norm_negative(model_limits):
    # Should raise if max_grad_norm < 0
    with pytest.raises(ValueError):
        create_finetune_request(
            model_limits=model_limits,
            training_file="train.jsonl",
            model="base-model",
            max_grad_norm=-1.0,
        ) # 6.35μs -> 6.17μs (2.88% faster)

def test_weight_decay_negative(model_limits):
    # Should raise if weight_decay < 0
    with pytest.raises(ValueError):
        create_finetune_request(
            model_limits=model_limits,
            training_file="train.jsonl",
            model="base-model",
            weight_decay=-0.01,
        ) # 6.26μs -> 6.80μs (8.01% slower)

def test_invalid_training_method(model_limits):
    # Should raise if training_method is not in AVAILABLE_TRAINING_METHODS
    with pytest.raises(ValueError):
        create_finetune_request(
            model_limits=model_limits,
            training_file="train.jsonl",
            model="base-model",
            training_method="invalid-method",
        ) # 6.50μs -> 6.54μs (0.734% slower)

def test_train_on_inputs_not_sft(model_limits):
    # Should raise if train_on_inputs is set but training_method is not "sft"
    with pytest.raises(ValueError):
        create_finetune_request(
            model_limits=model_limits,
            training_file="train.jsonl",
            model="base-model",
            training_method="dpo",
            train_on_inputs=True,
        ) # 5.41μs -> 5.25μs (3.16% faster)

def test_dpo_beta_not_dpo(model_limits):
    # Should raise if dpo_beta is set but training_method is not "dpo"
    with pytest.raises(ValueError):
        create_finetune_request(
            model_limits=model_limits,
            training_file="train.jsonl",
            model="base-model",
            dpo_beta=1.0,
        ) # 14.7μs -> 6.46μs (128% faster)

def test_dpo_normalize_logratios_by_length_not_dpo(model_limits):
    # Should raise if dpo_normalize_logratios_by_length is True but training_method is not "dpo"
    with pytest.raises(ValueError):
        create_finetune_request(
            model_limits=model_limits,
            training_file="train.jsonl",
            model="base-model",
            dpo_normalize_logratios_by_length=True,
        ) # 12.9μs -> 6.38μs (103% faster)

def test_rpo_alpha_not_dpo(model_limits):
    # Should raise if rpo_alpha is set but training_method is not "dpo"
    with pytest.raises(ValueError):
        create_finetune_request(
            model_limits=model_limits,
            training_file="train.jsonl",
            model="base-model",
            rpo_alpha=0.5,
        ) # 13.1μs -> 6.27μs (109% faster)

def test_rpo_alpha_negative(model_limits):
    # Should raise if rpo_alpha < 0
    with pytest.raises(ValueError):
        create_finetune_request(
            model_limits=model_limits,
            training_file="train.jsonl",
            model="base-model",
            training_method="dpo",
            rpo_alpha=-0.1,
        ) # 7.42μs -> 7.42μs (0.013% faster)

def test_simpo_gamma_not_dpo(model_limits):
    # Should raise if simpo_gamma is set but training_method is not "dpo"
    with pytest.raises(ValueError):
        create_finetune_request(
            model_limits=model_limits,
            training_file="train.jsonl",
            model="base-model",
            simpo_gamma=0.5,
        ) # 12.9μs -> 6.37μs (102% faster)

def test_simpo_gamma_negative(model_limits):
    # Should raise if simpo_gamma < 0
    with pytest.raises(ValueError):
        create_finetune_request(
            model_limits=model_limits,
            training_file="train.jsonl",
            model="base-model",
            training_method="dpo",
            simpo_gamma=-0.1,
        ) # 8.07μs -> 8.06μs (0.087% faster)

def test_lora_dropout_out_of_bounds(model_limits):
    # Should raise if lora_dropout < 0 or >= 1
    with pytest.raises(ValueError):
        create_finetune_request(
            model_limits=model_limits,
            training_file="train.jsonl",
            model="base-model",
            lora=True,
            lora_dropout=-0.01,
        ) # 4.91μs -> 1.78μs (176% faster)
    with pytest.raises(ValueError):
        create_finetune_request(
            model_limits=model_limits,
            training_file="train.jsonl",
            model="base-model",
            lora=True,
            lora_dropout=1.0,
        ) # 1.96μs -> 961ns (103% faster)

def test_scheduler_num_cycles_nonpositive(model_limits):
    # Should raise if scheduler_num_cycles <= 0 for cosine scheduler
    with pytest.raises(ValueError):
        create_finetune_request(
            model_limits=model_limits,
            training_file="train.jsonl",
            model="base-model",
            scheduler_num_cycles=0.0,
        ) # 15.1μs -> 9.14μs (64.7% faster)
    with pytest.raises(ValueError):
        create_finetune_request(
            model_limits=model_limits,
            training_file="train.jsonl",
            model="base-model",
            scheduler_num_cycles=-1.0,
        ) # 7.12μs -> 3.81μs (87.1% faster)

# 3. LARGE SCALE TEST CASES

def test_large_batch_size_max(model_limits):
    # Test with batch_size="max" for large batch sizes
    codeflash_output = create_finetune_request(
        model_limits=model_limits,
        training_file="train.jsonl",
        model="base-model",
        batch_size="max",
    ); req = codeflash_output # 34.0μs -> 28.5μs (19.4% faster)

def test_large_batch_size_within_limits(model_limits):
    # Test with batch_size just below the max
    codeflash_output = create_finetune_request(
        model_limits=model_limits,
        training_file="train.jsonl",
        model="base-model",
        batch_size=model_limits.full_training.max_batch_size,
    ); req = codeflash_output # 29.1μs -> 23.1μs (26.0% faster)

def test_large_batch_size_dpo(model_limits):
    # Test with batch_size just below max_batch_size_dpo for DPO
    codeflash_output = create_finetune_request(
        model_limits=model_limits,
        training_file="train.jsonl",
        model="base-model",
        batch_size=model_limits.full_training.max_batch_size_dpo,
        training_method="dpo",
    ); req = codeflash_output # 24.5μs -> 24.4μs (0.455% faster)

def test_large_lora_rank(model_limits):
    # Test with lora_r at max_rank
    codeflash_output = create_finetune_request(
        model_limits=model_limits,
        training_file="train.jsonl",
        model="base-model",
        lora=True,
        lora_r=model_limits.lora_training.max_rank,
    ); req = codeflash_output # 33.3μs -> 24.7μs (35.1% faster)

def test_many_epochs(model_limits):
    # Test with a large number of epochs
    codeflash_output = create_finetune_request(
        model_limits=model_limits,
        training_file="train.jsonl",
        model="base-model",
        n_epochs=999,
    ); req = codeflash_output # 28.0μs -> 21.6μs (29.7% faster)

def test_large_scale_all_fields(model_limits):
    # Test with all optional fields set and large values
    codeflash_output = create_finetune_request(
        model_limits=model_limits,
        training_file="train.jsonl",
        model="base-model",
        validation_file="val.jsonl",
        n_epochs=100,
        n_evals=50,
        n_checkpoints=20,
        batch_size=64,
        learning_rate=0.01,
        lr_scheduler_type="cosine",
        min_lr_ratio=0.9,
        scheduler_num_cycles=10,
        warmup_ratio=0.5,
        max_grad_norm=10.0,
        weight_decay=0.1,
        lora=True,
        lora_r=16,
        lora_dropout=0.2,
        lora_alpha=32,
        lora_trainable_modules="all-linear",
        suffix="experiment-large",
        wandb_api_key="api-key",
        wandb_base_url="http://wandb.example.com",
        wandb_project_name="project",
        wandb_name="run-1",
        train_on_inputs=True,
        training_method="sft",
        from_checkpoint=None,
        from_hf_model=None,
        hf_model_revision="rev1",
        hf_api_token="token",
        hf_output_repo_name="repo",
    ); req = codeflash_output # 25.7μs -> 24.5μs (4.78% faster)

def test_large_scale_dpo_with_all_fields(model_limits):
    # Test DPO with all relevant fields and large values
    codeflash_output = create_finetune_request(
        model_limits=model_limits,
        training_file="train.jsonl",
        model="base-model",
        batch_size=model_limits.full_training.max_batch_size_dpo,
        training_method="dpo",
        dpo_beta=1.5,
        dpo_normalize_logratios_by_length=True,
        rpo_alpha=0.5,
        simpo_gamma=0.8,
        n_epochs=50,
        n_evals=10,
        n_checkpoints=5,
    ); req = codeflash_output # 330μs -> 326μs (1.29% faster)

def test_large_scale_lora_batch(model_limits):
    # Test lora training with large batch size at max
    codeflash_output = create_finetune_request(
        model_limits=model_limits,
        training_file="train.jsonl",
        model="base-model",
        lora=True,
        batch_size=model_limits.lora_training.max_batch_size,
    ); req = codeflash_output # 36.9μs -> 26.7μs (38.0% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from together.resources.finetune import create_finetune_request

⏪ Replay Tests and Runtime

To edit these changes git checkout codeflash/optimize-create_finetune_request-mgzsj4q6 and push.