Skip to main content

core.trainers.ebft.async_trainer

core.trainers.ebft.async_trainer

Async EBFT on TRL’s AsyncGRPOTrainer.

The rollout process (EBFTRolloutLoop) generates and leaves the advantage pending; EBFTScorer scores each optimizer step’s groups on the trainer, collectively on every rank: a no-grad forward of each completion’s feature sequence and its group’s reference through the frozen feature network (the actor with adapters disabled, or a frozen copy for full fine-tuning), pooled hidden states at ebft_feature_layers, then group_feature_rewards. See axolotl.core.trainers.async_grpo.scoring.

Classes

Name Description
AxolotlAsyncEBFTConfig AxolotlAsyncGRPOConfig plus EBFT’s fields.
AxolotlAsyncEBFTTrainer Structured-mode EBFT on async GRPO (see module doc).
EBFTScorer EBFT’s feature-matching reward, scored on the trainer.

AxolotlAsyncEBFTConfig

core.trainers.ebft.async_trainer.AxolotlAsyncEBFTConfig(
    model_type=None,
    lr_quadratic_warmup=False,
    pretraining=False,
    sample_packing=False,
    batch_flattening=False,
    balance_labels=False,
    label_balance_window_optim_steps=1,
    label_keep_prob=None,
    label_keep_prob_normalize='dropped',
    sample_packing_sequentially=False,
    sample_packing_mp_start_method=None,
    sample_packing_drop_attention_mask=False,
    multipack_real_batches=False,
    checkpoint_save_retry_wait=None,
    checkpoint_save_retries=1,
    include_tkps=True,
    eval_sample_packing=None,
    sample_packing_efficiency=1.0,
    sample_packing_bin_size=200,
    sample_packing_group_size=100000,
    max_seq_length=2048,
    dataset_num_proc=None,
    relora_prune_ratio=None,
    relora_prune_method=None,
    jagged_restart_steps=None,
    jagged_restart_warmup_steps=None,
    jagged_restart_anneal_steps=None,
    bench_split='eval',
    bench_dataset='pharaouk/dharma-1/dharma_1_mini.json',
    do_bench_eval=False,
    do_causal_lm_eval=False,
    max_bench_samples=None,
    bench_source_max_len=2048,
    dataloader_prefetch_factor=None,
    cosine_min_lr_ratio=None,
    cosine_constant_lr_ratio=None,
    loraplus_lr_ratio=None,
    loraplus_lr_embedding=1e-06,
    embedding_lr_scale=None,
    lr_groups=None,
    embedding_lr=None,
    qlora=False,
    orpo_alpha=None,
    lisa_n_layers=None,
    lisa_step_interval=None,
    lisa_layers_attribute=None,
    curriculum_sampling=None,
    alternate_lr_scheduler_type=None,
    chat_template=None,
    adam_beta3=None,
    adam_epsilon2=None,
    activation_offloading=None,
    layer_offloading=None,
    image_size=None,
    image_resize_algorithm=None,
    dion_learning_rate=None,
    dion_momentum=None,
    dion_rank_fraction=None,
    dion_rank_multiple_of=None,
    reward_weights=None,
    multi_objective_aggregation='sum_then_normalize',
    scale_rewards='group',
    mask_truncated_completions=False,
    skip_zero_advantage_batches=True,
    reward_num_workers=1,
    scoring_token_budget=None,
    loss_type='dapo',
    vllm_importance_sampling_mode='token_truncate',
    vllm_importance_sampling_cap=3.0,
    rollout_loop=None,
    rollout_loop_kwargs=None,
    vllm_merge_check_interval=DEFAULT_MERGE_CHECK_INTERVAL,
    vllm_lora_alias=None,
    ebft_feature_layers=(lambda: [0.25, 0.5, 0.75])(),
    ebft_embed_method='last_token',
    ebft_use_whitening=False,
    ebft_alignment_coef=1.0,
    ebft_diversity_coef=1.0,
    ebft_ce_coef=0.0,
    ebft_adaptive_max_tokens=True,
    ebft_gt_length_multiplier=1.5,
    generation_kwargs=None,
)

AxolotlAsyncGRPOConfig plus EBFT’s fields.

AxolotlAsyncEBFTTrainer

core.trainers.ebft.async_trainer.AxolotlAsyncEBFTTrainer(
    model,
    args=None,
    train_dataset=None,
    processing_class=None,
    callbacks=None,
    optimizers=(None, None),
    peft_config=None,
    **kwargs,
)

Structured-mode EBFT on async GRPO (see module doc).

EBFTScorer

core.trainers.ebft.async_trainer.EBFTScorer(
    layer_indices,
    embed_method,
    use_whitening,
    alignment_coef,
    diversity_coef,
    feature_network=None,
)

EBFT’s feature-matching reward, scored on the trainer.