Buckets:
| import logging | |
| from pathlib import Path | |
| import torch | |
| from train.fsdp_utils import ( | |
| AdapterEMA, | |
| FSDPCheckpoint, | |
| FSDPConfig, | |
| ddp_options, | |
| materialize_meta_buffers, | |
| model_construction_dtype, | |
| prepare_replicated_lora, | |
| ) | |
| class _LoRAModel(torch.nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| self.base = torch.nn.Linear(2, 2, bias=False) | |
| self.lora_A = torch.nn.Linear(2, 1, bias=False) | |
| self.base.requires_grad_(False) | |
| class _Rotary(torch.nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| self.config = object() | |
| self.rope_kwargs = {} | |
| self.rope_init_fn = lambda _config, device, **_kwargs: ( | |
| torch.ones(3, device=device), | |
| 1.0, | |
| ) | |
| self.register_buffer("inv_freq", torch.empty(3, device="meta"), persistent=False) | |
| self.original_inv_freq = self.inv_freq | |
| def _ddp_config(): | |
| return FSDPConfig( | |
| sharding_strategy="FULL_SHARD", | |
| backward_prefetch="BACKWARD_PRE", | |
| cpu_offload=False, | |
| num_replicate=1, | |
| distributed_strategy="ddp", | |
| ) | |
| def test_ddp_model_construction_uses_bfloat16_and_restores_default(): | |
| previous = torch.get_default_dtype() | |
| with model_construction_dtype(_ddp_config()): | |
| assert torch.get_default_dtype() == torch.bfloat16 | |
| model = _LoRAModel() | |
| assert model.base.weight.dtype == torch.bfloat16 | |
| assert model.base.weight.is_meta | |
| assert torch.get_default_dtype() == previous | |
| prepare_replicated_lora(model) | |
| assert model.base.weight.dtype == torch.bfloat16 | |
| assert model.lora_A.weight.dtype == torch.float32 | |
| def test_adapter_ema_updates_swaps_and_restores_only_lora(): | |
| model = _LoRAModel() | |
| original_base = model.base.weight.detach().clone() | |
| original_adapter = model.lora_A.weight.detach().clone() | |
| ema = AdapterEMA(model) | |
| with torch.no_grad(): | |
| model.lora_A.weight.add_(2) | |
| ema.update(model, decay=0.5) | |
| expected_ema = original_adapter + 1 | |
| trained_adapter = model.lora_A.weight.detach().clone() | |
| model.train() | |
| with ema.apply(model) as validation_model: | |
| assert not validation_model.training | |
| assert torch.allclose(validation_model.lora_A.weight, expected_ema) | |
| assert torch.equal(validation_model.base.weight, original_base) | |
| assert model.training | |
| assert torch.equal(model.lora_A.weight, trained_adapter) | |
| def test_replicated_lora_keeps_only_adapter_master_weights_in_fp32(): | |
| model = prepare_replicated_lora(_LoRAModel()) | |
| assert model.base.weight.dtype == torch.bfloat16 | |
| assert model.lora_A.weight.dtype == torch.float32 | |
| def test_cpu_ema_matches_resident_ema_and_restores_student(): | |
| model = _LoRAModel() | |
| resident = AdapterEMA(model) | |
| offloaded = AdapterEMA(model, cpu_offload=True) | |
| for _ in range(3): | |
| with torch.no_grad(): | |
| model.lora_A.weight.add_(0.125) | |
| resident.update(model, 0.9999) | |
| offloaded.update(model, 0.9999) | |
| for name, value in resident.parameters.items(): | |
| assert offloaded.parameters[name].device.type == "cpu" | |
| assert torch.equal(value, offloaded.parameters[name]) | |
| student = model.lora_A.weight.detach().clone() | |
| with offloaded.apply(model): | |
| assert torch.equal(model.lora_A.weight, resident.parameters['lora_A.weight']) | |
| assert torch.equal(model.lora_A.weight, student) | |
| def test_lora_ddp_uses_dynamic_unused_parameter_detection(): | |
| options = ddp_options(2) | |
| assert options["find_unused_parameters"] is True | |
| assert options["static_graph"] is False | |
| assert options["broadcast_buffers"] is True | |
| assert options["device_ids"] == [2] | |
| def test_nonpersistent_rotary_buffer_is_materialized_after_meta_load(): | |
| rotary = _Rotary() | |
| materialize_meta_buffers(rotary) | |
| assert not rotary.inv_freq.is_meta | |
| assert torch.equal(rotary.inv_freq, torch.ones(3)) | |
| assert rotary.original_inv_freq is rotary.inv_freq | |
| def test_ddp_adapter_checkpoint_is_atomic_and_resumable(tmp_path, monkeypatch): | |
| monkeypatch.setattr("train.fsdp_utils.dist.get_rank", lambda: 0) | |
| monkeypatch.setattr("train.fsdp_utils.dist.get_world_size", lambda: 1) | |
| monkeypatch.setattr("train.fsdp_utils.dist.barrier", lambda: None) | |
| monkeypatch.setattr( | |
| "train.fsdp_utils.dist.broadcast_object_list", lambda *_args, **_kwargs: None | |
| ) | |
| monkeypatch.setattr( | |
| "train.fsdp_utils.torch.cuda.get_rng_state", lambda: torch.zeros(1) | |
| ) | |
| monkeypatch.setattr( | |
| "train.fsdp_utils.torch.cuda.set_rng_state", lambda *_args: None | |
| ) | |
| model = _LoRAModel() | |
| ema = AdapterEMA(model) | |
| optimizer = torch.optim.AdamW(model.lora_A.parameters(), lr=1e-3) | |
| scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lambda _step: 1.0) | |
| model.lora_A.weight.sum().backward() | |
| optimizer.step() | |
| scheduler.step() | |
| checkpoint = Path( | |
| FSDPCheckpoint.fsdp_save_ckpt( | |
| ckpt_dir=tmp_path, | |
| train_steps=7, | |
| model=model, | |
| ema_model=ema, | |
| optimizer=optimizer, | |
| scheduler=scheduler, | |
| data_status=[{"logical_train": {0: [0, 7, 3]}}], | |
| logger=logging.getLogger("ddp-checkpoint-test"), | |
| fsdp_config=_ddp_config(), | |
| adapter_only=True, | |
| adapter_metadata={"rank": 128}, | |
| ) | |
| ) | |
| assert (checkpoint / "COMPLETE").is_file() | |
| assert (checkpoint / "adapter.safetensors").is_file() | |
| assert (checkpoint / "ema_adapter.safetensors").is_file() | |
| assert (checkpoint / "optimizer.00000-of-00001.pt").is_file() | |
| optimizer, scheduler, step, status = FSDPCheckpoint.try_load_train_state( | |
| checkpoint, optimizer, scheduler, _ddp_config() | |
| ) | |
| assert step == 8 | |
| assert status == {"logical_train": {0: [0, 7, 3]}} | |
Xet Storage Details
- Size:
- 5.82 kB
- Xet hash:
- 1443a637c19b2fcd5439c2e7a15ff1da4f020907594028ba1e306208ee9c107e
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.