(self)
| 36 | world_size = 2 |
| 37 | |
| 38 | def test_model(self): |
| 39 | lora_config = LoRAConfig(lora_r=16, lora_alpha=16, base_weight_sharding=2) |
| 40 | quant_config = None |
| 41 | hidden_dim = 64 |
| 42 | nlayers = 4 |
| 43 | |
| 44 | with deepspeed.linear.Init(lora_config=lora_config, quant_config=quant_config): |
| 45 | model = SimpleModel(hidden_dim=hidden_dim, nlayers=nlayers) |
| 46 | |
| 47 | init_lora(model) |
| 48 | |
| 49 | model_norms = [model.linears[i].weight.norm().item() for i in range(nlayers)] |
| 50 | |
| 51 | ds_config = { |
| 52 | "train_batch_size": 2, |
| 53 | "steps_per_print": 1, |
| 54 | "bf16": { |
| 55 | "enabled": True |
| 56 | }, |
| 57 | "optimizer": { |
| 58 | "type": "Adam", |
| 59 | "params": { |
| 60 | "lr": 0.00015 |
| 61 | } |
| 62 | }, |
| 63 | "zero_optimization": { |
| 64 | "stage": 1 |
| 65 | } |
| 66 | } |
| 67 | model, *_ = deepspeed.initialize(config=ds_config, model=model, model_parameters=model.parameters()) |
| 68 | |
| 69 | engine_norms = [model.module.linears[i].weight.norm().item() for i in range(nlayers)] |
| 70 | |
| 71 | # Ensure that sharded weights are not broadcast during engine init |
| 72 | assert engine_norms == model_norms, f"{dist.get_rank()=} base weight norms are not the same after engine init, {engine_norms=} != {model_norms=}" |
| 73 | |
| 74 | data_loader = random_dataloader(model=model, |
| 75 | total_samples=50, |
| 76 | hidden_dim=hidden_dim, |
| 77 | device=model.device, |
| 78 | dtype=torch.bfloat16) |
| 79 | for n, batch in enumerate(data_loader): |
| 80 | loss = model(batch[0], batch[1]) |
| 81 | model.backward(loss) |
| 82 | model.step() |
| 83 | |
| 84 | |
| 85 | @pytest.mark.skip( |
nothing calls this directly
no test coverage detected