[BUG] AttributeError: 'Linear' object has no attribute 'ds_grads_remaining'

[jianrui2025]

Describe the bug
I met this issue by using huggingface'trainer, accelerate, peft and deepspeed. when I set modules_to_save=["lm_head"] in LoraConfig, this issue occured. when I set modules_to_save=None in LoraConfig, code run successful. when I set deepspeed=0.16.4, this code will run succesful with modules_to_save=["lm_head"]

model:Qwen/Qwen3-1.7B
dataset:krisfu/delicate_medical_r1_data_chinese
deepspeed=0.17.6

To Reproduce

import json
import pandas as pd
import torch
from datasets import Dataset
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer, DataCollatorForSeq2Seq
import os
from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
from peft import LoraConfig, get_peft_model,PeftModel
from accelerate import Accelerator



# import swanlab

if __name__ == "__main__":
    tokenizer = AutoTokenizer.from_pretrained("/home/jianrui/workspace/ChinChunMei-LLM/demo/check_transformers/Qwen/Qwen3-1.7B")
    model = AutoModelForCausalLM.from_pretrained("/home/jianrui/workspace/ChinChunMei-LLM/demo/check_transformers/Qwen/Qwen3-1.7B")

    config = LoraConfig(
        r=16,
        lora_alpha=16,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
        lora_dropout=0.1,
        bias="none",
        modules_to_save=["lm_head"],
    )
    
    model = get_peft_model(model, config)


    model.print_trainable_parameters()
    model.enable_input_require_grads()

    def process_func(example):
        """
        将数据集进行预处理，tokenization操作。
        """ 
        PROMPT = "你是一个医学专家，你需要根据用户的问题，给出带有思考的回答。"
        MAX_LENGTH = 1000
        input_ids, attention_mask, labels = [], [], []
        instruction = tokenizer(
            f"<|im_start|>system\n{PROMPT}<|im_end|>\n<|im_start|>user\n{example['input']}<|im_end|>\n<|im_start|>assistant\n",
            add_special_tokens=False,
        )
        response = tokenizer(f"{example['output']}", add_special_tokens=False)
        input_ids = instruction["input_ids"] + response["input_ids"] + [tokenizer.pad_token_id]
        attention_mask = instruction["attention_mask"] + response["attention_mask"] + [1]
        labels = [-100] * len(instruction["input_ids"]) + response["input_ids"] + [tokenizer.pad_token_id]
        if len(input_ids) > MAX_LENGTH:  # 做一个截断
            input_ids = input_ids[:MAX_LENGTH]
            attention_mask = attention_mask[:MAX_LENGTH]
            labels = labels[:MAX_LENGTH]
        return {"input_ids": input_ids, "attention_mask": attention_mask, "labels": labels}
    train_dataset = load_dataset("json", data_files="/home/jianrui/workspace/ChinChunMei-LLM/demo/check_transformers/krisfu/delicate_medical_r1_data_chinese/r1_data_example_instruct.jsonl",split="train")
    train_dataset = train_dataset.map(process_func)


    args = TrainingArguments(
        output_dir="/home/jianrui/workspace/ChinChunMei-LLM/demo/check_transformers/Qwen/Qwen3-1.7B_output",
        per_device_train_batch_size=2,  
        gradient_accumulation_steps=1,
        # 验证操作
        per_device_eval_batch_size=1,
        eval_strategy="no",
        eval_steps=100,
        logging_steps=1,
        num_train_epochs=12,
        # 模型保存. 模型保存16位模型。
        save_strategy="steps",
        save_steps=100,
        save_total_limit =5,
        save_safetensors =True,
        save_only_model =False, 

        learning_rate=1e-4,
        save_on_each_node=False,
        gradient_checkpointing=True,
        # report_to="swanlab",
        run_name="qwen3-1.7B",
        fp16=True,
        warmup_ratio=0.1,
        data_seed=38,
        # log_level="info",
    )

    trainer = Trainer(
        model=model,
        args=args,
        train_dataset=train_dataset,
        data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer, padding=True),
    )

    trainer.train()

Expected behavior
I want this code running successful with modules_to_save=["lm_head"].

ds_report output
Please run ds_report to give us details about your setup.

Screenshots
If applicable, add screenshots to help explain your problem.

System info (please complete the following information):

• OS: [e.g. Ubuntu 18.04]
• GPU count and types [e.g. two machines with x8 A100s each]
• Interconnects (if applicable) [e.g., two machines connected with 100 Gbps IB]
• Python version
• Any other relevant info about your setup

Launcher context
Are you launching your experiment with the deepspeed launcher, MPI, or something else?

Docker context
Are you using a specific docker image that you can share?

Additional context
Add any other context about the problem here.

[Flink-ddd]

Hi @jianrui2025 , Could you provide more information about the model file you have the problem with?

[jianrui2025]

model file? for example?

[lipf1024]

same problem, For example, in Qwen2.5-VL-3B, its embedding and lm_head are shared. When I use PEFT to perform LoRA training and set --lora_modules_to_save embed_tokens, a similar problem occurs.

[zihaohe123]

Same problem here

[Functionhx]

I reproduced this issue on DeepSpeed v0.17.6 with ZeRO-3 and PEFT modules_to_save=["lm_head"]. The same minimal repro passes on v0.16.4 and when modules_to_save=None.

I'd like to work on a fix and add a regression test.