ai-agents-fsdp-activation-chec-when-using-fsdp-with-activation-checkpointing-true-3dbad06a

THEN Set `use_cache=False` in the model keyword arguments when activation checkpointing is enabled. Change the condition from `use_cache=not gradient_checkpointing` to `use_cache=not (gradient_checkpointing or activation_checkpointing)`. For example: `model_kwargs = {'use_cache': not (sft_config.gradient_checkpointing or sft_config.fsdp_config.activation_checkpointing)}`.

model_kwargs = dict(
    attn_implementation=sft_config.attn_implementation,
    torch_dtype=sft_config.torch_dtype,
    use_cache=not (sft_config.gradient_checkpointing or sft_config.fsdp_config.activation_checkpointing)
)
model = AutoModelForCausalLM.from_pretrained(sft_config.model_name_or_path, **model_kwargs)