fix ppo callbacks

[leuitong]

This fix restores user callbacks on the PPO trainer that TRL's PPOTrainer.init silently discards.

PR type

• [✅ ] Bug Fix
• New Feature
• Document Updates
• More Models or Datasets Support

PR information

Problem

Swift's user-trainer callbacks (e.g. early_stop, lisa, adalora, activation_cpu_offload, perf_log or other user-define callbacks) were silently dropped and never fired during PPO training.

Root cause: SwiftMixin.init registers these callbacks on the HF Trainer by passing callbacks=... into HFTrainer.init (swift/trainers/mixin.py). But the swift PPOTrainer then invokes TRL's PPOTrainer.init (ppo_trainer_init) with a new_kwargs whitelist that
did not include callbacks. TRL's init then runs:

self.callbacks = default_callbacks if callbacks is None else default_callbacks + callbacks
self.callback_handler = CallbackHandler(self.callbacks, ...)

Because callbacks is None, TRL rebuilds callback_handler from scratch with only default/integration callbacks, discarding every callback the SwiftMixin had just registered. Consequently, features relying on these callbacks (early stopping, LISA layer swapping,
activation CPU offload, perf logging) were inactive for PPO runs.

Fix

In swift/rlhf_trainers/ppo_trainer.py:

1. Removed 'callbacks' from the new_kwargs whitelist so TRL rebuilds a clean handler.
2. After ppo_trainer_init returns, re-register the swift callbacks onto the newly built handler:

for callback in self.args.callbacks:
self.add_callback(callbacks_map[callback](self.args, self))

This mirrors SwiftMixin._get_callbacks exactly (instantiate callbacks_map[name](args, trainer)), sharing the same construction paths/signatures across swift versions. Placing it after TRL's init attaches to the new callback_handler instead of the discarded one. At
this point the handler holds only default/reporting callbacks, so re-adding causes no dedup warnings and no double-invocation (the originals died with the old handler).

Experiment results

Paste your experiment result here(if needed).

[2026-08-07 19:30:20] run sh: `/usr/local/bin/python -m torch.distributed.run --nproc_per_node 1 --master_port 23456 --nnodes 1 --node_rank 0 --master_addr aistudio-nc9lot71-ptjob-master-0 /usr/local/lib/python3.11/site-packages/swift/cli/rlhf.py --template qwen3 --tuner_type lora --log_completions true --dataloader_num_workers 1 --use_vllm true --gradient_accumulation_steps 4 --offload_model true --top_p 0.85 --offload_optimizer true --num_train_epochs 2 --lr_scheduler_type cosine --num_generations 4 --save_steps 100 --temperature 1 --sleep_level 1 --lora_merge False --max_length 4096 --gc_collect_after_offload true --model_type qwen3 --warmup_ratio 0.1 --vllm_gpu_memory_utilization 0.5 --max_completion_length 4096 --vllm_mode colocate --torch_dtype bfloat16 --vllm_max_model_len 4096 --per_device_train_batch_size 1 --learning_rate 5e-06 --logging_steps 10 --overlong_filter true --finetuning_type lora --overwrite_cache true --overwrite_output_dir true --stage ppo --trust_remote_code true --use_fast_tokenizer true --rlhf_type ppo --output_dir /checkpoint --model /rl-mnt/model/actor_model --dataset /data/ds_0 --report_to tensorboard --logging_dir /home/admin/logs/tfevent --callbacks custom_post_training_callback --ignore_args_error True --reward_model /rl-mnt/model/reward_model`

A custom callback was specified but did not take effect in the logs. --callbacks custom_post_training_callback

{'eps': 0, 'objective/kl': '3.323', 'objective/entropy': '2083', 'objective/non_score_reward': '-0.1661', 'objective/rlhf_reward': '2.631', 'objective/scores': '2.797', 'policy/approxkl_avg': '0.0006469', 'policy/clipfrac_avg': '0.001334', 'loss/policy_avg': '-0.002931', 'loss/value_avg': '1.114', 'val/clipfrac_avg': '0.01949', 'policy/entropy_avg': '0.6767', 'val/ratio': '0.9998', 'val/ratio_var': '8e-08', 'val/num_eos_tokens': 3, 'lr': '5e-06', 'episode': 16, 'epoch': '1.6', 'global_step/max_steps': '2/3', 'elapsed_time': '6m 42s', 'remaining_time': '3m 21s', 'memory(GiB)': '114.2', 'train_speed(s/it)': '201.1'}
[2026-08-06 17:16:00] 
Train:  33%|███▎      | 1/3 [06:42<06:43, 201.99s/it]
Train:  33%|███▎      | 1/3 [06:42<06:43, 201.99s/it]
Train:  67%|██████▋   | 2/3 [06:42<03:21, 201.00s/it]2026-08-06 17:16:00,340 - accelerate.accelerator - INFO -
                                                  
{'eps': 0, 'objective/kl': '3.431', 'objective/entropy': '1856', 'objective/non_score_reward': '-0.1715', 'objective/rlhf_reward': '3.75', 'objective/scores': '3.922', 'policy/approxkl_avg': '0.0006374', 'policy/clipfrac_avg': '0.001667', 'loss/policy_avg': '-0.001021', 'loss/value_avg': '1.103', 'val/clipfrac_avg': '0.006758', 'policy/entropy_avg': '0.6533', 'val/ratio': '1', 'val/ratio_var': '4.2e-07', 'val/num_eos_tokens': 4, 'lr': '2.5e-06', 'episode': 24, 'epoch': '2.4', 'global_step/max_steps': '3/3', 'elapsed_time': '10m 4s', 'remaining_time': '0s', 'memory(GiB)': '114.2', 'train_speed(s/it)': '201.3'}

Train:  67%|██████▋   | 2/3 [10:03<03:21, 201.00s/it]
Train:  67%|██████▋   | 2/3 [10:03<03:21, 201.00s/it]
Train: 100%|██████████| 3/3 [10:03<00:00, 201.28s/it][INFO:swift] Saving model checkpoint to /mnt/ais/checkpoint/0000000002/179494994/1/v0-20260806-170439/checkpoint-3

Train: 100%|██████████| 3/3 [10:04<00:00, 201.51s/it]
[INFO:swift] last_model_checkpoint: /mnt/ais/checkpoint/0000000002/179494994/1/v0-20260806-170439/checkpoint-3
[INFO:swift] best_model_checkpoint: None
[INFO:swift] images_dir: /mnt/ais/checkpoint/0000000002/179494994/1/v0-20260806-170439/images
[2026-08-06 17:16:01] 2026-08-06 17:16:01,126 - matplotlib.font_manager - INFO - generated new fontManager
[2026-08-06 17:16:02] [INFO:swift] End time of running main: 2026-08-06 17:16:02.901463
[rank0]:[W806 17:16:03.662950460 ProcessGroupNCCL.cpp:1553] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator())

now fixed：

2026-08-07 12:21:37,521 - AisPostTrainingSwiftCallback - INFO - basic train config：
2026-08-07 12:21:37,521 - AisPostTrainingSwiftCallback - INFO - - total batch size: 4
2026-08-07 12:21:37,521 - AisPostTrainingSwiftCallback - INFO - - train epoch: 2.0
Passing `generation_config` together with generation-related arguments=({'output_scores', 'return_dict_in_generate'}) is deprecated and will be removed in future versions. Please pass either a `generation_config` object OR all generation parameters explicitly, but not both.
{'eps': 0, 'objective/kl': '1.973', 'objective/entropy': '1092', 'objective/non_score_reward': '-0.09864', 'objective/rlhf_reward': '3.058', 'objective/scores': '3.156', 'policy/approxkl_avg': '0.0005298', 'policy/clipfrac_avg': '0.001672', 'loss/policy_avg': '-0.002344', 'loss/value_avg': '3.488', 'val/clipfrac_avg': '0', 'policy/entropy_avg': '0.6152', 'val/ratio': '1', 'val/ratio_var': '2.8e-07', 'val/num_eos_tokens': 3, 'lr': '0', 'episode': 4, 'epoch': '0.4', 'global_step/max_steps': '1/5', 'elapsed_time': '3m 41s', 'remaining_time': '14m 43s', 'memory(GiB)': '47.83', 'train_speed(s/it)': '220.6'}
Train:  20%|██        | 1/5 [03:40<14:42, 220.64s/it]2026-08-07 12:29:05,835 - accelerate.accelerator - INFO - [RANK 0] The used dataset had no length, returning gathered tensors. You should drop the remainder yourself.
{'eps': 0, 'objective/kl': '2.392', 'objective/entropy': '3025', 'objective/non_score_reward': '-0.1196', 'objective/rlhf_reward': '1.169', 'objective/scores': '1.289', 'policy/approxkl_avg': '0.0008044', 'policy/clipfrac_avg': '0.001609', 'loss/policy_avg': '-0.001554', 'loss/value_avg': '2.602', 'val/clipfrac_avg': '0.08885', 'policy/entropy_avg': '0.7797', 'val/ratio': '0.9999', 'val/ratio_var': '3.2e-07', 'val/num_eos_tokens': 1, 'lr': '2.5e-06', 'episode': 16, 'epoch': '1.6', 'global_step/max_steps': '4/5', 'elapsed_time': '14m 56s', 'remaining_time': '3m 44s', 'memory(GiB)': '47.85', 'train_speed(s/it)': '224.1'}
{'eps': 0, 'objective/kl': '4.25', 'objective/entropy': '1273', 'objective/non_score_reward': '-0.2125', 'objective/rlhf_reward': '2.116', 'objective/scores': '2.328', 'policy/approxkl_avg': '0.0005599', 'policy/clipfrac_avg': '0.001413', 'loss/policy_avg': '-0.006166', 'loss/value_avg': '2.818', 'val/clipfrac_avg': '0.0005732', 'policy/entropy_avg': '0.682', 'val/ratio': '0.9997', 'val/ratio_var': '1.8e-07', 'val/num_eos_tokens': 2, 'lr': '7.3e-07', 'episode': 20, 'epoch': '2', 'global_step/max_steps': '5/5', 'elapsed_time': '18m 37s', 'remaining_time': '0s', 'memory(GiB)': '47.85', 'train_speed(s/it)': '223.4'}
Train: 100%|██████████| 5/5 [18:36<00:00, 222.87s/it][INFO:swift] Saving model checkpoint to /checkpoint/0000000002/179515115/1/v0-20260807-122120/checkpoint-5
2026-08-07 12:40:15,059 - AisPostTrainingSwiftCallback - INFO - === save ckpt ===
2026-08-07 12:40:15,059 - AisPostTrainingSwiftCallback - INFO - step: 5
2026-08-07 12:40:15,059 - AisPostTrainingSwiftCallback - INFO - current loss: N/A
2026-08-07 12:40:15,059 - AisPostTrainingSwiftCallback - INFO - best loss: inf
2026-08-07 12:40:15,059 - AisPostTrainingSwiftCallback - INFO - save dir: /mnt/ais/checkpoint/0000000002/179515115/1/v0-20260807-122120
2026-08-07 12:40:15,654 - AisPostTrainingSwiftCallback - INFO - === train end ===