能否增加对qwen3的支持？或者如何在其他框架中使用（ada）lomo优化器？

[LRC-bit]

我曾经尝试在trl框架中使用lomo优化器，但是总是会出现代码报错的问题。
于是我根据lomo优化器项目下的推荐来尝试CoLLiE框架，但是我发现这个框架有阵子没更新了。
我想知道是否有可能比较简单地增加Qwen3模型的支持呢？或者是否可能在trl中正常跑通多卡zero3训练的例子？
期待您的回复。

[KaiLv69]

你好~

1. 请问有trl具体的报错信息吗，我可以提个pr修复
2. 对于qwen3这种collie自己没有实现的模型，可以直接使用transformers库的AutoModel传给trainer训练。
   

   CoLLiE/collie/controller/trainer.py

   Line 59 in 3826b3e

   * transformers 提供的模型 ``transformers.PreTrainedModel`` 只支持 `ZeRO`

你可以尝试修改一下这个例子，把模型换成transformers里的qwen，去掉evaluator这些东西。https://github.com/OpenMOSS/CoLLiE/blob/main/examples/one_sentence_overfitting/finetune_moss_for_training.py

[LRC-bit]

trl在单卡训练的时候是正常的。在使用accelerate+deepspeed zero3的时候会出现问题。我凭印象先简单阐述一下，adalomo.py文件下：
for n, p in self.model.named_parameters():
if self.zero3_enabled:
if len(p.ds_shape) == 1:
self.exp_avg_sq[n] = torch.zeros(p.ds_shape[0], dtype=torch.float32).cuda()
else:
self.exp_avg_sq_row[n] = torch.zeros(p.ds_shape[0], dtype=torch.float32).cuda()
self.exp_avg_sq_col[n] = torch.zeros(p.ds_shape[1], dtype=torch.float32).cuda()
此处的p.ds_shape会报错没有ds_shape这个属性。我看到最新的release说是修复了这个问题，但是好像没有改善。
当我把此处的p.ds_shape简单修改成p.shape测试的时候会出现另外一个报错。提示的是deepspeed中的某一个函数（具体忘了）接收到了预料之外的learning_rate参数。
感谢，请问您有什么头绪吗~具体的报错信息如果必要的话我想我可以再复现一次

[LRC-bit]

另外，我还有个问题是，adalomo优化器的学习率相比传统的adamw有简单的经验倍率关系吗？例如在adamw中我使用1e-4的学习率可以得到很好的收敛效果，那么在adalomo中我应该设置为多少较为合适呢？

[KaiLv69]

1. 这个ds_shape报错我确实是修过一次。可以给一个最小复现的代码吗，我可以再试试
2. adalomo学习率可以调的比adamw大一点

[LRC-bit]

from transformers import AutoTokenizer, AutoModelForCausalLM
from datasets import load_dataset
from trl import SFTTrainer,SFTConfig

def main(input_model_dir, input_dataset, output_dir, lora_enable=False):
    tokenizer = AutoTokenizer.from_pretrained(input_model_dir)
    tokenizer.padding_side = 'left'  

    dataset = load_dataset('json', data_files=input_dataset, split='train')

    model = AutoModelForCausalLM.from_pretrained(  
        input_model_dir,
        torch_dtype = 'bfloat16',   
        use_cache = False,  
    )

    training_args = SFTConfig(
        bf16=True,  
        per_device_train_batch_size=128,
        gradient_accumulation_steps=1, 
        optim='adalomo',  

        logging_steps=10, 
        save_steps=2000,  
        output_dir=output_dir,
        report_to="none",
    ) 

    trainer = SFTTrainer(
        model=model, 
        processing_class=tokenizer, 
        args=training_args,   
        train_dataset=dataset, 
    )

    train_result = trainer.train()
    metrics = train_result.metrics
    metrics["train_samples"] = len(dataset)
    trainer.log_metrics("train", metrics)
    trainer.save_metrics("train", metrics)
    trainer.save_state()
    trainer.save_model(training_args.output_dir)

if __name__ == "__main__":
    input_model_dir = 'Qwen3-4B'
    input_dataset = 'xxx.jsonl'
    output_dir = 
    main(input_model_dir, input_dataset, output_dir, lora_enable)

以上是训练代码，使用accelerate launch train.py运行
下面是accelerate的配置文件config.yaml

compute_environment: LOCAL_MACHINE
debug: false
deepspeed_config:
  gradient_accumulation_steps: 1
  offload_optimizer_device: none
  offload_param_device: none
  zero3_init_flag: true
  zero3_save_16bit_model: true
  zero_stage: 3
distributed_type: DEEPSPEED
downcast_bf16: 'no'
dynamo_config:
  dynamo_backend: INDUCTOR
enable_cpu_affinity: false
machine_rank: 0
main_training_function: main
mixed_precision: bf16
num_machines: 1
num_processes: 2
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false

运行之后会得到如下报错：

[rank1]: Traceback (most recent call last):
[rank1]:   File "/home/linruichen/lomoLLM/trl_train_copy.py", line 48, in <module>
[rank1]:     main(input_model_dir, input_dataset, output_dir)
[rank1]:   File "/home/linruichen/lomoLLM/trl_train_copy.py", line 36, in main
[rank1]:     train_result = trainer.train()
[rank1]:   File "/home/linruichen/miniconda3/envs/lomoLLM/lib/python3.10/site-packages/transformers/trainer.py", line 2240, in train
[rank1]:     return inner_training_loop(
[rank1]:   File "/home/linruichen/miniconda3/envs/lomoLLM/lib/python3.10/site-packages/transformers/trainer.py", line 2322, in _inner_training_loop
[rank1]:     self.optimizer, self.lr_scheduler = deepspeed_init(self, num_training_steps=max_steps)
[rank1]:   File "/home/linruichen/miniconda3/envs/lomoLLM/lib/python3.10/site-packages/transformers/integrations/deepspeed.py", line 467, in deepspeed_init
[rank1]:     optimizer, lr_scheduler = deepspeed_optim_sched(
[rank1]:   File "/home/linruichen/miniconda3/envs/lomoLLM/lib/python3.10/site-packages/transformers/integrations/deepspeed.py", line 387, in deepspeed_optim_sched
[rank1]:     optimizer = trainer.create_optimizer()
[rank1]:   File "/home/linruichen/miniconda3/envs/lomoLLM/lib/python3.10/site-packages/transformers/trainer.py", line 1246, in create_optimizer
[rank1]:     self.optimizer = optimizer_cls(optimizer_grouped_parameters, **optimizer_kwargs)
[rank1]:   File "/home/linruichen/miniconda3/envs/lomoLLM/lib/python3.10/site-packages/lomo_optim/adalomo.py", line 86, in __init__
[rank1]:     if len(p.ds_shape) == 1:
[rank1]: AttributeError: 'Parameter' object has no attribute 'ds_shape'. Did you mean: 'is_shared'?

我不清楚这里的ds_shape指的是什么，所以凭感觉将adalomo.py中的ds_shape修改成shape，获得的新报错为：

[rank1]: Traceback (most recent call last):
[rank1]:   File "/home/linruichen/lomoLLM/trl_train_copy.py", line 48, in <module>
[rank1]:     main(input_model_dir, input_dataset, output_dir)
[rank1]:   File "/home/linruichen/lomoLLM/trl_train_copy.py", line 36, in main
[rank1]:     train_result = trainer.train()
[rank1]:   File "/home/linruichen/miniconda3/envs/lomoLLM/lib/python3.10/site-packages/transformers/trainer.py", line 2240, in train
[rank1]:     return inner_training_loop(
[rank1]:   File "/home/linruichen/miniconda3/envs/lomoLLM/lib/python3.10/site-packages/transformers/trainer.py", line 2555, in _inner_training_loop
[rank1]:     tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
[rank1]:   File "/home/linruichen/miniconda3/envs/lomoLLM/lib/python3.10/site-packages/trl/trainer/sft_trainer.py", line 864, in training_step
[rank1]:     return super().training_step(*args, **kwargs)
[rank1]:   File "/home/linruichen/miniconda3/envs/lomoLLM/lib/python3.10/site-packages/transformers/trainer.py", line 3791, in training_step
[rank1]:     self.accelerator.backward(loss, **kwargs)
[rank1]:   File "/home/linruichen/miniconda3/envs/lomoLLM/lib/python3.10/site-packages/accelerate/accelerator.py", line 2545, in backward
[rank1]:     self.deepspeed_engine_wrapped.backward(loss, sync_gradients=self.sync_gradients, **kwargs)
[rank1]:   File "/home/linruichen/miniconda3/envs/lomoLLM/lib/python3.10/site-packages/accelerate/utils/deepspeed.py", line 270, in backward
[rank1]:     self.engine.backward(loss, **kwargs)
[rank1]:   File "/home/linruichen/miniconda3/envs/lomoLLM/lib/python3.10/site-packages/deepspeed/utils/nvtx.py", line 20, in wrapped_fn
[rank1]:     ret_val = func(*args, **kwargs)
[rank1]: TypeError: DeepSpeedEngine.backward() got an unexpected keyword argument 'learning_rate'

目前就是这么个情况，我挺好奇问题的原因是什么，希望得到您的解答。另外您提到adalomo的学习率可以调得比adamw大一些，能更具体一点吗？比如是倍数关系还是常数关系呢？

[KaiLv69]

在使用deepspeed的时候p.ds_shape是因为用zero3后，parameter的p.shape不一定准。这里报没有ds_shape可能是1.accelerate又包了一层或者2.这里还没有使用deepspeed的zero3切模型。

下面那个新报错你可以看看这个pr的lomo_backward(self, loss: torch.Tensor, learning_rate: float)：https://github.com/huggingface/accelerate/pull/2695/files

我不确定backward有没有调用这个函数，如果没有调用这个函数，确实会出现下面的报错。

关于学习率，可以先试试1e-3

[LRC-bit]

我刚刚从1e-2试到了5e-5，发现确实1e-3的效果是最佳的