Buckets:
| import logging | |
| from types import SimpleNamespace | |
| import pytest | |
| import torch | |
| from safetensors.torch import save_file | |
| from train import pretrain_unified_navit as pretrain | |
| from train.fsdp_utils import FSDPCheckpoint | |
| class _Position(torch.nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| self.pos_embed = torch.nn.Parameter(torch.zeros(2), requires_grad=False) | |
| class _Model(torch.nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| self.weight = torch.nn.Parameter(torch.zeros(2)) | |
| self.latent_pos_embed = _Position() | |
| self.vit_pos_embed = _Position() | |
| def test_hf_ema_model_only_load_allows_only_regenerated_positions(tmp_path): | |
| payload = { | |
| "weight": torch.ones(2), | |
| "latent_pos_embed.pos_embed": torch.ones(2), | |
| "vit_pos_embed.pos_embed": torch.ones(2), | |
| } | |
| save_file(payload, tmp_path / "ema.safetensors") | |
| model, ema = FSDPCheckpoint.try_load_ckpt( | |
| tmp_path, | |
| logging.getLogger("test"), | |
| _Model(), | |
| _Model(), | |
| resume_from_ema=True, | |
| ) | |
| assert torch.equal(model.weight, torch.ones(2)) | |
| assert torch.equal(ema.weight, torch.ones(2)) | |
| assert torch.equal(model.latent_pos_embed.pos_embed, torch.zeros(2)) | |
| def test_ema_initialization_falls_back_to_thinkmorph_model_payload(tmp_path): | |
| payload = { | |
| "weight": torch.ones(2), | |
| "latent_pos_embed.pos_embed": torch.ones(2), | |
| "vit_pos_embed.pos_embed": torch.ones(2), | |
| } | |
| save_file(payload, tmp_path / "model.safetensors") | |
| model, _ = FSDPCheckpoint.try_load_ckpt( | |
| tmp_path, | |
| logging.getLogger("test"), | |
| _Model(), | |
| resume_from_ema=True, | |
| ) | |
| assert torch.equal(model.weight, torch.ones(2)) | |
| def test_ddp_meta_model_assigns_checkpoint_without_random_materialization(tmp_path): | |
| payload = { | |
| "weight": torch.ones(2, dtype=torch.bfloat16), | |
| "latent_pos_embed.pos_embed": torch.ones(2, dtype=torch.bfloat16), | |
| "vit_pos_embed.pos_embed": torch.ones(2, dtype=torch.bfloat16), | |
| } | |
| save_file(payload, tmp_path / "ema.safetensors") | |
| with torch.device("meta"): | |
| model = _Model() | |
| model, _ = FSDPCheckpoint.try_load_ckpt( | |
| tmp_path, | |
| logging.getLogger("test"), | |
| model, | |
| resume_from_ema=True, | |
| assign=True, | |
| ) | |
| assert not model.weight.is_meta | |
| assert model.weight.dtype == torch.bfloat16 | |
| assert torch.equal(model.latent_pos_embed.pos_embed, torch.ones(2)) | |
| def test_main_keeps_sequentially_loaded_ema_model(monkeypatch, tmp_path): | |
| """The model-only checkpoint API returns the loaded model in slot zero.""" | |
| class _Tokenizer: | |
| def __len__(self): | |
| return 32 | |
| class _ReachedEmaSetup(Exception): | |
| pass | |
| model_args = pretrain.ModelArguments(model_path=str(tmp_path)) | |
| data_args = pretrain.DataArguments() | |
| training_args = pretrain.TrainingArguments( | |
| results_dir=str(tmp_path / "logs"), | |
| checkpoint_dir=str(tmp_path / "checkpoints"), | |
| resume_from=str(tmp_path), | |
| resume_model_only=True, | |
| finetune_from_ema=True, | |
| finetune_from_hf=True, | |
| total_steps=1, | |
| stop_after_steps=1, | |
| gradient_accumulation_steps=1, | |
| log_every=1, | |
| save_every=1, | |
| sharding_strategy="FULL_SHARD", | |
| num_shard=4, | |
| ) | |
| parser = SimpleNamespace( | |
| parse_args_into_dataclasses=lambda: (model_args, data_args, training_args) | |
| ) | |
| monkeypatch.setattr(pretrain, "HfArgumentParser", lambda *_: parser) | |
| monkeypatch.setattr(pretrain.torch.cuda, "is_available", lambda: True) | |
| monkeypatch.setattr(pretrain.torch.cuda, "device_count", lambda: 4) | |
| monkeypatch.setattr(pretrain.torch.cuda, "set_device", lambda *_: None) | |
| monkeypatch.setattr(pretrain.dist, "init_process_group", lambda *_: None) | |
| monkeypatch.setattr(pretrain.dist, "get_rank", lambda: 1) | |
| monkeypatch.setattr(pretrain.dist, "get_world_size", lambda: 4) | |
| monkeypatch.setattr(pretrain.dist, "barrier", lambda: None) | |
| monkeypatch.setattr(pretrain, "create_logger", lambda *_: logging.getLogger("test")) | |
| monkeypatch.setattr(pretrain, "set_seed", lambda *_: None) | |
| monkeypatch.setattr( | |
| pretrain.Qwen2Config, | |
| "from_json_file", | |
| staticmethod(lambda *_: SimpleNamespace()), | |
| ) | |
| monkeypatch.setattr( | |
| pretrain.SiglipVisionConfig, | |
| "from_json_file", | |
| staticmethod(lambda *_: SimpleNamespace(num_hidden_layers=2)), | |
| ) | |
| monkeypatch.setattr( | |
| pretrain, | |
| "load_ae", | |
| lambda **_: (torch.nn.Identity(), SimpleNamespace()), | |
| ) | |
| monkeypatch.setattr( | |
| pretrain.Qwen2Tokenizer, | |
| "from_pretrained", | |
| staticmethod(lambda *_: _Tokenizer()), | |
| ) | |
| monkeypatch.setattr( | |
| pretrain, | |
| "add_special_tokens", | |
| lambda tokenizer: (tokenizer, [], 0), | |
| ) | |
| monkeypatch.setattr( | |
| pretrain, | |
| "build_bagel_model", | |
| lambda *_: torch.nn.Linear(1, 1), | |
| ) | |
| monkeypatch.setattr( | |
| pretrain.FSDPCheckpoint, | |
| "try_load_ckpt", | |
| staticmethod( | |
| lambda _path, _logger, model, ema_model=None, resume_from_ema=False: ( | |
| model, | |
| ema_model, | |
| ) | |
| ), | |
| ) | |
| monkeypatch.setattr(pretrain, "distributed_wrapper", lambda model, *_: model) | |
| def _capture_ema(_student, model, _config): | |
| assert model is not None | |
| raise _ReachedEmaSetup | |
| monkeypatch.setattr(pretrain, "distributed_ema_setup", _capture_ema) | |
| with pytest.raises(_ReachedEmaSetup): | |
| pretrain.main() | |
Xet Storage Details
- Size:
- 5.62 kB
- Xet hash:
- b3f9b895858c28059185f85490d89b9e00efb4bf7cda8b81f736514536202e76
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.