fsdp

def create_training_arguments(epochs, batch_size, learning_rate, output_dir):
    training_args = TrainingArguments(
        output_dir=output_dir,
        num_train_epochs=epochs,
        per_device_train_batch_size=batch_size,
        learning_rate=learning_rate,
        save_strategy="no",
        do_eval = True,
        evaluation_strategy='epoch',
        log_level_replica='passive',
    )
    return training_args