| """Decoupled DPO backward == coupled backward, to numerical identity. |
| |
| The 7B stage backwards each side's trunk separately (both checkpointed |
| graphs resident at once measured 21.5GB and OOM'd). The decoupling must be |
| a pure memory transformation: same loss, same metrics, same parameter |
| gradients as the coupled path the kernel tests already certify. |
| """ |
|
|
| import pytest |
| import torch |
|
|
| from align.model.mistral import NmoMistralForCausalLM |
| from align.train.dpo_loss import DpoLossConfig, make_dpo_loss_adapter |
|
|
| IGNORE_INDEX = -100 |
|
|
|
|
| @pytest.fixture |
| def tiny(tiny_model_constants): |
| from test_trainer import tiny_config |
|
|
| return tiny_config(tiny_model_constants) |
|
|
|
|
| def _batch(vocab: int, b: int = 2, t: int = 12, seed: int = 7) -> dict: |
| g = torch.Generator().manual_seed(seed) |
| ids_c = torch.randint(0, vocab, (b, t), generator=g) |
| ids_r = torch.randint(0, vocab, (b, t), generator=g) |
| lab_c = ids_c.clone() |
| lab_r = ids_r.clone() |
| lab_c[:, : t // 2] = IGNORE_INDEX |
| lab_r[:, : t // 2] = IGNORE_INDEX |
| return { |
| "input_ids_chosen": ids_c, |
| "labels_chosen": lab_c, |
| "input_ids_rejected": ids_r, |
| "labels_rejected": lab_r, |
| "ref_logp_chosen": torch.randn(b, generator=g), |
| "ref_logp_rejected": torch.randn(b, generator=g), |
| } |
|
|
|
|
| def _grads(model) -> dict: |
| return { |
| n: p.grad.detach().clone() |
| for n, p in model.named_parameters() |
| if p.grad is not None and p.grad.abs().sum() > 0 |
| } |
|
|
|
|
| def test_decoupled_backward_matches_coupled(tiny): |
| torch.manual_seed(0) |
| model = NmoMistralForCausalLM(tiny, device="cpu") |
| batch = _batch(tiny.vocab_size) |
|
|
| coupled = make_dpo_loss_adapter(DpoLossConfig(), require_fused=True) |
| out = coupled(model, batch) |
| out.loss.backward() |
| g_coupled = _grads(model) |
| assert g_coupled, "coupled path produced no gradients" |
| loss_coupled = float(out.loss.detach()) |
| model.zero_grad(set_to_none=True) |
|
|
| dec = make_dpo_loss_adapter( |
| DpoLossConfig(), require_fused=True, decoupled_backward=True, grad_scale=1.0 |
| ) |
| out2 = dec(model, batch) |
| |
| g_before_proxy = _grads(model) |
| out2.loss.backward() |
| g_after_proxy = _grads(model) |
| assert set(g_before_proxy) == set(g_after_proxy) |
| for n in g_before_proxy: |
| assert torch.equal(g_before_proxy[n], g_after_proxy[n]), ( |
| f"proxy backward mutated model grad {n}" |
| ) |
|
|
| assert abs(float(out2.loss.detach()) - loss_coupled) < 1e-6 |
| assert set(out2.metrics) == set(out.metrics) |
| assert set(g_after_proxy) == set(g_coupled) |
| for n in g_coupled: |
| assert torch.allclose(g_coupled[n], g_after_proxy[n], rtol=1e-5, atol=1e-7), n |
|
|
|
|
| def test_decoupled_grad_scale_scales_real_grads(tiny): |
| torch.manual_seed(0) |
| model = NmoMistralForCausalLM(tiny, device="cpu") |
| batch = _batch(tiny.vocab_size) |
| dec1 = make_dpo_loss_adapter( |
| DpoLossConfig(), require_fused=True, decoupled_backward=True, grad_scale=1.0 |
| ) |
| dec1(model, batch) |
| g1 = _grads(model) |
| model.zero_grad(set_to_none=True) |
| dec4 = make_dpo_loss_adapter( |
| DpoLossConfig(), require_fused=True, decoupled_backward=True, grad_scale=0.25 |
| ) |
| dec4(model, batch) |
| g4 = _grads(model) |
| for n in g1: |
| assert torch.allclose(g1[n] * 0.25, g4[n], rtol=1e-5, atol=1e-8), n |
|
|
