Buckets:
| import copy | |
| import pytest | |
| import torch | |
| from shared.reference import FrozenReference | |
| class _TinyActor(torch.nn.Module): | |
| def __init__(self) -> None: | |
| super().__init__() | |
| self.live = torch.nn.Linear(2, 1, bias=False) | |
| self.frozen = torch.nn.Parameter( | |
| torch.tensor([3.0]), requires_grad=False | |
| ) | |
| def forward(self, values: torch.Tensor) -> torch.Tensor: | |
| return self.live(values) + self.frozen | |
| def test_reference_context_is_no_grad_and_restores_changed_actor( | |
| ) -> None: | |
| actor = _TinyActor() | |
| actor.live.weight.data.copy_(torch.tensor([[1.0, 2.0]])) | |
| reference = FrozenReference(actor) | |
| actor.live.weight.data.copy_(torch.tensor([[5.0, 7.0]])) | |
| actor.frozen.data.fill_(13.0) | |
| actor.train() | |
| with reference.apply(actor) as reference_actor: | |
| output = reference_actor(torch.tensor([[2.0, 3.0]])) | |
| assert torch.equal(output, torch.tensor([[21.0]])) | |
| assert output.requires_grad is False | |
| assert reference_actor.training is False | |
| assert torch.equal(actor.live.weight, torch.tensor([[5.0, 7.0]])) | |
| assert torch.equal(actor.frozen, torch.tensor([13.0])) | |
| assert actor.training is True | |
| def test_reference_state_describes_an_fp32_immutable_native_snapshot() -> None: | |
| actor = _TinyActor().to(torch.bfloat16) | |
| reference = FrozenReference(actor) | |
| state = reference.state_dict() | |
| descriptor = reference.descriptor() | |
| assert state["schema"] == "shared_frozen_reference_v1" | |
| assert state["update_invariant"] == "never" | |
| assert state["parameter_names"] == ["live.weight"] | |
| assert state["parameter_specs"]["live.weight"] == { | |
| "shape": [1, 2], | |
| "native_dtype": "bfloat16", | |
| } | |
| assert state["shadow"]["live.weight"].dtype == torch.float32 | |
| assert len(state["descriptor_hash"]) == 64 | |
| assert descriptor["descriptor_hash"] == state["descriptor_hash"] | |
| assert not hasattr(reference, "update") | |
| def test_resume_restores_run_start_reference_into_an_already_changed_actor() -> None: | |
| initial_actor = _TinyActor() | |
| initial_actor.live.weight.data.copy_(torch.tensor([[1.0, 2.0]])) | |
| state = FrozenReference(initial_actor).state_dict() | |
| resumed_actor = _TinyActor() | |
| resumed_actor.live.weight.data.copy_(torch.tensor([[9.0, 10.0]])) | |
| restored = FrozenReference(resumed_actor) | |
| restored.load_state_dict(state) | |
| assert torch.equal(resumed_actor.live.weight, torch.tensor([[9.0, 10.0]])) | |
| with restored.apply(resumed_actor) as reference_actor: | |
| assert torch.equal( | |
| reference_actor(torch.tensor([[2.0, 3.0]])), | |
| torch.tensor([[11.0]]), | |
| ) | |
| assert restored.descriptor() == { | |
| key: value for key, value in state.items() if key != "shadow" | |
| } | |
| def test_reference_state_cannot_be_reloaded_while_the_view_is_active() -> None: | |
| actor = _TinyActor() | |
| reference = FrozenReference(actor) | |
| state = reference.state_dict() | |
| with reference.apply(actor): | |
| with pytest.raises(RuntimeError, match="active"): | |
| reference.load_state_dict(state) | |
| def test_reference_rejects_a_changed_trainable_parameter_schema() -> None: | |
| actor = _TinyActor() | |
| reference = FrozenReference(actor) | |
| state = reference.state_dict() | |
| actor.frozen.requires_grad_(True) | |
| with pytest.raises(RuntimeError, match="schema changed"): | |
| reference.load_state_dict(state) | |
| def test_reference_resume_rejects_corrupt_state_and_hashes() -> None: | |
| actor = _TinyActor() | |
| reference = FrozenReference(actor) | |
| state = reference.state_dict() | |
| corrupt_snapshot = copy.deepcopy(state) | |
| corrupt_snapshot["shadow"]["live.weight"][0, 0] += 1.0 | |
| with pytest.raises(ValueError, match="snapshot hash"): | |
| reference.load_state_dict(corrupt_snapshot) | |
| corrupt_descriptor = copy.deepcopy(state) | |
| corrupt_descriptor["descriptor_hash"] = "0" * 64 | |
| with pytest.raises(ValueError, match="descriptor hash"): | |
| reference.load_state_dict(corrupt_descriptor) | |
| corrupt_invariant = copy.deepcopy(state) | |
| corrupt_invariant["update_invariant"] = "optimizer_step" | |
| with pytest.raises(ValueError, match="never-update invariant"): | |
| reference.load_state_dict(corrupt_invariant) | |
| def test_reference_rejects_a_nonfinite_run_start_actor() -> None: | |
| actor = _TinyActor() | |
| actor.live.weight.data.fill_(float("nan")) | |
| with pytest.raises(ValueError, match="non-finite"): | |
| FrozenReference(actor) | |
Xet Storage Details
- Size:
- 4.46 kB
- Xet hash:
- c4c5359defe8aabd2c2abc8f6b8d39834c2873ef98b3887a1cb3eba08c5ac582
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.