use_gradient_checkpointing in FastLanguageModel.get_peft_model not working.

[sutakori]

Issue

After the following code, the model.is_gradient_checkpointing is False and True when before and after get_peft_model, where I want False.

use_gradient_checkpointing = False

model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=model_name,
        max_seq_length=4096,
        dtype=torch.bfloat16,
        load_in_4bit=False,
        device_map={"": f"cuda:{os.environ.get('LOCAL_RANK', '0')}"},
        use_gradient_checkpointing=use_gradient_checkpointing,
        use_cache=True,
)

model = FastLanguageModel.get_peft_model(
        model,
        r=16,
        target_modules=['q_proj', 'k_proj', 'v_proj', 'o_proj', 'up_proj', 'down_proj', 'gate_proj'],
        lora_dropout=0,
        bias='none',
        use_gradient_checkpointing=use_gradient_checkpointing,
)

Debugging

FastLanguageModel.get_peft_model calls prepare_model_for_training in unsloth_zoo, for line 199-200 in unsloth_zoo/training_utils.py

if hasattr(model, "_set_gradient_checkpointing"):
       model._set_gradient_checkpointing()

the model._set_gradient_checkpointing() is transformers library take "enable = True" as default, thus alter the model.is_gradient_checkpointing