import torch, deepspeed
from deepspeed.runtime.zero import unwrap_model_for_generation
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import LoraConfig, get_peft_model, TaskType

m = "hf-internal-testing/tiny-random-gpt2"
tok = AutoTokenizer.from_pretrained(m); tok.pad_token = tok.eos_token
model = get_peft_model(AutoModelForCausalLM.from_pretrained(m, dtype=torch.bfloat16),
                       LoraConfig(task_type=TaskType.CAUSAL_LM, r=4, target_modules=["c_attn"]))
model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})

cfg = {"train_micro_batch_size_per_gpu": 1, "bf16": {"enabled": True},
       "zero_optimization": {"stage": 3, "stage3_param_persistence_threshold": 0},
       "optimizer": {"type": "Adam", "params": {"lr": 1e-3}}}
engine, *_ = deepspeed.initialize(model=model, config=cfg,
                                  model_parameters=[p for p in model.parameters() if p.requires_grad])

ids = tok("hello", return_tensors="pt").input_ids.to(engine.device)
for _ in range(2):
    with unwrap_model_for_generation(engine) as unwrapped:
        with torch.no_grad():
            unwrapped.generate(ids, max_new_tokens=4, do_sample=False, pad_token_id=tok.pad_token_id)
    out = engine(input_ids=ids, labels=ids)
    engine.backward(out.loss); engine.step()