diff --git a/.gitignore b/.gitignore index 03e62300..0222862f 100644 --- a/.gitignore +++ b/.gitignore @@ -149,4 +149,5 @@ log *.qdrep .vscode -!bmtrain/dist \ No newline at end of file +!bmtrain/dist +tests/test_log.txt \ No newline at end of file diff --git a/README-ZH.md b/README-ZH.md index 264edbd5..598777b2 100644 --- a/README-ZH.md +++ b/README-ZH.md @@ -273,19 +273,16 @@ BMTrain 支持**所有** PyTorch 原生的优化器和损失函数，同时你 ### 第 4 部分: 训练 ```python -# 新建损失缩放器实例 -loss_scaler = bmtrain.optim.LossScaler() -# 将所有的 optimzer 及(可选)其对应的 lr_scheduler 收入损失缩放器管理。 -loss_scaler.add_optimizer(optimizer, lr_scheduler) +# 新建优化器管理器实例 +optim_manager = bmtrain.optim.OptimManager(loss_scale=1024) +# 将所有的 optimzer 及(可选)其对应的 lr_scheduler 收入优化器管理器管理。 +optim_manager.add_optimizer(optimizer, lr_scheduler) # 可以再次调用 add_optimizer 加入其他优化器 for iteration in range(1000): # ... 为每个rank加载数据 ... - # 梯度清零 - loss_scaler.zero_grad() # 为每个 optimizer 调用 zero_grad - - # 前向传播 + # 前向传播并计算梯度 pos = torch.arange(enc_input.size(1)).long().cuda().repeat(enc_input.size(0), 1) logits = model( enc_input, @@ -296,14 +293,19 @@ for iteration in range(1000): loss = loss_func(logits.view(batch * seq_len, vocab_out_size), targets.view(batch * seq_len)) - global_loss = bmtrain.sum_loss(loss).item() # 聚合所有rank上的损失 + global_loss = bmtrain.sum_loss(loss).item() # 聚合所有rank上的损失, 仅用于输出训练日志 + + # 梯度清零 + optim_manager.zero_grad() # 为每个 optimizer 调用 zero_grad # 损失缩放和反向传播 - loss = loss_scaler(loss) - loss.backward() + optim_manager.backward(loss) + + # 梯度裁剪 + grad_norm = optim_manager.clip_grad_norm(optimizer.param_groups, max_norm=1.0) # 更新参数 - loss_scaler.step() + optim_manager.step() # ... 保存checkpoint、打印日志 ... ``` @@ -312,40 +314,11 @@ for iteration in range(1000): 你可以根据代码中的注释来了解各部分代码的作用。 -唯一需要说明的是 `loss_scaler`，**损失缩放**是混合精度训练中的一项常用技术，需要在反向传播前通过 `loss_scaler(loss)` 对 `loss` 进行放缩，用于避免梯度下溢。在使用损失所放后，优化器的 `step()` 等步骤需要有一些细节上的调整。我们在 `loss_scaler` 帮你实现了这些细节, 你只需要通过 `add_optimizer` 将优化器和学习率调整策略收入 `loss_scaler` 管理，并由 `loss_scaler` 代为执行 `zero_grad()` 和 `step()` 操作。 - -如果你没有使用 BMTrain 中的融合优化器，你可以不用 `loss_scaler`, 相应的代码修改为: - -```python -for iteration in range(1000): - # ... 为每个rank加载数据 ... - - # 梯度清零 - optimizer.zero_grad() - - # 前向传播 - pos = torch.arange(enc_input.size(1)).long().cuda().repeat(enc_input.size(0), 1) - logits = model( - enc_input, - pos, - pos < enc_length[:, None] - ) - batch, seq_len, vocab_out_size = logits.size() - - loss = loss_func(logits.view(batch * seq_len, vocab_out_size), targets.view(batch * seq_len)) - - global_loss = bmtrain.sum_loss(loss).item() # 聚合所有rank上的损失 +唯一需要说明的是 `optim_manager`。在使用 BMTrain 后，优化器的部分相关操作需要有一些细节上的调整。我们在 `optim_manager` 帮你实现了这些细节, 你只需要通过 `add_optimizer` 将优化器和学习率调整策略收入 `optim_manager` 管理，并由 `optim_manger` 代为执行 `zero_grad()`, `backward()`, `clip_grad_norm()` 和 `step()` 等操作。 - # 反向传播 - loss.backward() - - # 更新参数 - bmtrain.optim_step(optimizer, lr_scheduler) - - # ... 保存checkpoint、打印日志 ... -``` +如果你没有使用混合精度训练，你可以不用损失缩放，只需要将 `OptimManger(loss_scale=None)` 构造函数中 `loss_scale` 置为 None 即可, 这也是 `OptimManager` 的默认构造参数。 -需要说明的是最后更新参数时，需要使用 `bmtrain.optim_step`, 而不能直接使用 `optimizer.step()` 和 `lr_scheduler.step()`。 +如果你使用了混合精度训练，**损失缩放**是混合精度训练中的一项常用技术，我们在 `optim_manager.backward(loss)` 帮你对 `loss` 进行了放缩，用于避免梯度下溢。只需要将 `OptimManger` 构造函数中 `loss_scale` 置为一个浮点数即可。 `loss_scale` 会在训练过程中根据梯度进行自适应的调整。
diff --git a/README.md b/README.md index 16e50037..45ed273c 100644 --- a/README.md +++ b/README.md @@ -276,19 +276,16 @@ In addition, BMTrain also provides the common LRScheduler in the `bmtrain.lr_sch ### Part 4: Training Loop ```python -# create a new instance of loss scaler -loss_scaler = bmtrain.optim.LossScaler() -# let loss_scaler handle all the optimizer and (optional) their corresponding lr_scheduler -loss_scaler.add_optimizer(optimizer, lr_scheduler) +# create a new instance of optimizer manager +optim_manager = bmtrain.optim.OptimManager() +# let optim_manager handle all the optimizer and (optional) their corresponding lr_scheduler +optim_manager.add_optimizer(optimizer, lr_scheduler) # add_optimizer can be called multiple times to add other optimizers. for iteration in range(1000): # ... load data for each rank ... - # zero grad - loss_scaler.zero_grad() # calling zero_grad for each optimizer - - # forward + # forward pass and calculate loss pos = torch.arange(enc_input.size(1)).long().cuda().repeat(enc_input.size(0), 1) logits = model( enc_input, @@ -299,14 +296,19 @@ for iteration in range(1000): loss = loss_func(logits.view(batch * seq_len, vocab_out_size), targets.view(batch * seq_len)) - global_loss = bmtrain.sum_loss(loss).item() # sum the loss across all ranks + global_loss = bmtrain.sum_loss(loss).item() # sum the loss across all ranks. This is only used for the training log + + # zero grad + optim_manager.zero_grad() # calling zero_grad for each optimizer + + # clip grad norm + grad_norm = optim_manager.clip_grad_norm(optimizer.param_groups, max_norm=1.0) # loss scale and backward - loss = loss_scaler(loss) - loss.backward() + optim_manager.backward() # optimizer step - loss_scaler.step() + optim_manager.step() # ... save checkpoint or print logs ... ``` @@ -315,9 +317,12 @@ The training loop part will be slightly longer, but just like a normal training You can follow the comments in the code to get an idea of what each section of code is doing. -The only additional note is `loss_scaler`, *loss scale* is the technique widely used in mixed precision training to prevent gradient underflow by adding `loss_scaler(loss)` to scale the `loss` before backward. When using loss scaling, some details in optimizers should be adjusted. We have implemented all those details needed in `loss_scaler`. What you need is just letting `loss_scaler` to handle all the optimizers by `add_optimizer`, and letting `loss_scaler` do `zero_grad()` and `step()` instead. +The only additional note is `optimizer`. After using BMTrain, some details in optimizers should be adjusted. We have implemented all those details needed in `optim_manager`. What you need is just letting `optim_manager` to handle all the optimizers by `add_optimizer`, and letting `optim_manager` do `zero_grad()`, `backward()`, `clip_grad_norm()` and `step()` instead. + +If you are not using the mixed-precision training, you can train without `loss_scale`. Just set `loss_scale` to None in the `__init__` function of `OptimManager(loss_scale=None)`, which is also the default. + +If you are using mixed-precision training, *loss scale* is the technique widely used in mixed precision training to prevent gradient underflow. By using `optim_manager.backward(loss)` to scale the `loss` before backward and set `loss_scale` to some floating number in the `__init__` function of `OptimManager`。The `loss_scale` would be adjusted adaptively based on the gradient during training. -If you are not using the fused optimizer in BMTrain, you can train without `loss_scaler`. The code will be changed as below: ```python for iteration in range(1000): diff --git a/bmtrain/__init__.py b/bmtrain/__init__.py index 25c72052..14ff50de 100644 --- a/bmtrain/__init__.py +++ b/bmtrain/__init__.py @@ -8,7 +8,6 @@ from .synchronize import synchronize, sum_loss, wait_loader, gather_result from .checkpointing import checkpoint from .block_layer import CheckpointBlock, TransformerBlockList -from .backward import optim_step from .wrapper import BMTrainModelWrapper from .pipe_layer import PipelineTransformerBlockList from . import debug diff --git a/bmtrain/backward.py b/bmtrain/backward.py deleted file mode 100644 index c503c972..00000000 --- a/bmtrain/backward.py +++ /dev/null @@ -1,28 +0,0 @@ -from typing import Optional -import torch -from .global_var import config -from .utils import print_rank -from .lr_scheduler.warmup import WarmupLRScheduler - - -def optim_step(optim : torch.optim.Optimizer, lr_scheduler : Optional[WarmupLRScheduler] = None): - """ - This is a helper function to call optimizer.step() and lr_scheduler.step() and synchronize streams. - - Args: - optim (torch.optim.Optimizer): A pytorch optimizer, e.g. torch.optim.Adam, torch.optim.SGD or bmtrain.optim.AdamOffloadOptimizer - lr_scheduler (Optional[WarmupLRScheduler]): A warmup lr scheduler, e.g. bmt.lr_scheduler.Noam - - This function can also handle gradient overflow by reducing the loss scale when it occurs. - - """ - - current_stream = torch.cuda.current_stream() - # some reduce ops of distributed parameter were launched on load stream - current_stream.wait_stream(config['load_stream']) - - optim.step() - if lr_scheduler is not None: - lr_scheduler.step() - - config['load_stream'].wait_stream(current_stream) \ No newline at end of file diff --git a/bmtrain/inspect/tensor.py b/bmtrain/inspect/tensor.py index 7d56ad7f..526ab8fc 100644 --- a/bmtrain/inspect/tensor.py +++ b/bmtrain/inspect/tensor.py @@ -129,9 +129,9 @@ def get_summary(self): nccl.allReduce( info.storage(), info.storage(), - "avg", + "sum", comm - ) + ) / nccl.commCount(comm) x_mean = info[0].cpu().item() x_std = math.sqrt(info[1].cpu().item()) grad_mean = None @@ -146,9 +146,9 @@ def get_summary(self): nccl.allReduce( info.storage(), info.storage(), - "avg", + "sum", comm - ) + ) / nccl.commCount(comm) x_mean = info[0].cpu().item() x_std = math.sqrt(info[1].cpu().item()) grad_mean = info[2].cpu().item() diff --git a/bmtrain/optim/__init__.py b/bmtrain/optim/__init__.py index 327d6676..15206328 100644 --- a/bmtrain/optim/__init__.py +++ b/bmtrain/optim/__init__.py @@ -1,3 +1,3 @@ from .adam import AdamOptimizer from .adam_offload import AdamOffloadOptimizer -from .loss_scaler import LossScaler \ No newline at end of file +from .optim_manager import OptimManager \ No newline at end of file diff --git a/bmtrain/optim/loss_scaler.py b/bmtrain/optim/optim_manager.py similarity index 76% rename from bmtrain/optim/loss_scaler.py rename to bmtrain/optim/optim_manager.py index 2a243863..e72fcddb 100644 --- a/bmtrain/optim/loss_scaler.py +++ b/bmtrain/optim/optim_manager.py @@ -27,20 +27,37 @@ def grad_rescale(param_groups, scale): if p.grad is not None and p.requires_grad: p.grad /= scale -class LossScaler: - """Loss scaler for mix-precision training +class OptimManager: + """wait cuda stream. Optional: add loss scaler for mix-precision training Args: - loss_scale (float): The initial loss scale. + loss_scale (float): The initial loss scale. Default to None for not using loss scaling. loss_scale_factor (float): The loss scale factor. loss_scale_steps (int): The loss scale steps. + + Examples: + >>> optim_manager = bmt.optim.OptimManager(loss_scale=1024) + >>> optim_manager.add_optimizer(optimizer1) + >>> optim_manager.add_optimizer(optimizer2, lr_scheduler2) + >>> for data in dataset: + >>> # forward pass and calculate loss + >>> optim_manager.zero_grad() + >>> optim_manager.backward(loss) + >>> optim_manager.clip_grad_norm(optimizer1.param_groups, max_norm=1.0, norm_type=2) + >>> optim_manager.clip_grad_norm(optimizer2.param_groups, max_norm=2.0, norm_type=2) + >>> optim_manager.step() """ def __init__(self, - loss_scale : float = 1, + loss_scale : Optional[float] = None, loss_scale_factor : float = 2, loss_scale_steps : int = 1024, ): - self.loss_scale = loss_scale + if loss_scale is not None: + self.loss_scale = loss_scale + self.loss_scale_enabled = True + else: + self.loss_scale = 1 + self.loss_scale_enabled = False self.steps_since_last_scale = 0 self.loss_scale_factor = loss_scale_factor if loss_scale_factor > 1 else 1 / loss_scale_factor self.loss_scale_steps = loss_scale_steps @@ -53,8 +70,8 @@ def add_optimizer( optimizer: torch.optim.Optimizer, lr_scheduler: Optional[WarmupLRScheduler] = None, ): - """Add optimizer and (optional) its corresponding lr_scheduler into loss_scaler. - All optimizers in the same loss_scaler share the same loss scale. + """Add optimizer and (optional) its corresponding lr_scheduler into optim_manager. + All optimizers in the same optim_manager share the same loss scale. Args: optim (torch.optim.Optimizer): A pytorch optimizer, e.g. torch.optim.Adam, torch.optim.SGD or bmtrain.optim.AdamOffloadOptimizer @@ -63,14 +80,21 @@ def add_optimizer( self.optimizers.append(optimizer) self.lr_schedulers.append(lr_scheduler) - def __call__(self, loss : torch.Tensor) -> torch.Tensor: + def loss_scale(self, loss : torch.Tensor) -> torch.Tensor: + return loss * (self.loss_scale / config['world_size']) # loss scale + + def backward(self, loss : torch.Tensor): """ Backward with loss scale. Args: loss (torch.Tensor): loss """ - return loss * (self.loss_scale * config["pipe_size"] / config['world_size']) + loss = self.loss_scale(loss) + loss.backward() + # some reduce ops of distributed parameter were launched on load stream + current_stream = torch.cuda.current_stream() + current_stream.wait_stream(config['load_stream']) def zero_grad(self): """ @@ -88,11 +112,7 @@ def step(self): This function can also handle gradient overflow by reducing the loss scale when it occurs. """ - current_stream = torch.cuda.current_stream() - # some reduce ops of distributed parameter were launched on load stream - current_stream.wait_stream(config['load_stream']) - - if self.loss_scale > 1: + if self.loss_scale_enabled and self.loss_scale > 1: has_overflow = False for optimizer in self.optimizers: try: @@ -110,17 +130,20 @@ def step(self): if hasattr(optimizer, "_bmtrain_optimizer") and optimizer._bmtrain_optimizer: optimizer.step(scale=self.loss_scale) else: - grad_rescale(optimizer.param_groups, self.loss_scale) + if self.loss_scale_enabled: + grad_rescale(optimizer.param_groups, self.loss_scale) optimizer.step() if lr_scheduler is not None: lr_scheduler.step() - self.steps_since_last_scale += 1 + if self.loss_scale_enabled: + self.steps_since_last_scale += 1 - if self.steps_since_last_scale >= self.loss_scale_steps: - self._justify_scale(self.loss_scale * self.loss_scale_factor) + if self.steps_since_last_scale >= self.loss_scale_steps: + self._justify_scale(self.loss_scale * self.loss_scale_factor) + current_stream = torch.cuda.current_stream() config['load_stream'].wait_stream(current_stream) def clip_grad_norm(self, param_groups, max_norm, norm_type=2, eps=1e-6): @@ -136,17 +159,8 @@ def clip_grad_norm(self, param_groups, max_norm, norm_type=2, eps=1e-6): Returns: Total norm of the parameters (viewed as a single vector). - - Examples: - >>> optimizer = bmt.optim.AdamOffloadOptimizer(model.parameters()) - >>> loss_scaler = bmt.optim.LossScaler() - >>> loss_scaler.add_optimizer(optimizer) - >>> # ... - >>> # backward_step() - >>> loss_scaler.clip_grad_norm(optimizer.param_groups, max_norm=1.0, norm_type=2) - """ - scale = self.loss_scale * config['pipe_size'] / config['world_size'] + scale = self.loss_scale grads = [] parameters = [p for group in param_groups for p in group['params']] for p in parameters: diff --git a/bmtrain/param_init.py b/bmtrain/param_init.py index c95f90f8..125283e3 100644 --- a/bmtrain/param_init.py +++ b/bmtrain/param_init.py @@ -22,7 +22,6 @@ def init_distributed_parameter(params : Iterable[torch.nn.Parameter]): param._init_method(tmp_tensor) # Pytorch 1.11 changed the API of storage.__getitem__ - # use zero_rank to support pipeline torch.tensor([], dtype=param.dtype, device=param.device).set_(param.storage())[:] = \ torch.tensor([], dtype=param.dtype, device=param.device).set_(tmp_storage)[partition_size * config['rank'] : partition_size * (config['rank'] + 1)] # param.storage().copy_(tmp_storage[partition_size * config['rank'] : partition_size * (config['rank'] + 1)]) diff --git a/bmtrain/synchronize.py b/bmtrain/synchronize.py index f704d159..80d39057 100644 --- a/bmtrain/synchronize.py +++ b/bmtrain/synchronize.py @@ -31,7 +31,7 @@ def sum_loss(loss : torch.Tensor): This is a helper function to reduce the loss across all workers. """ warnings.warn("bmtrain.sum_loss is deprecated and will be removed in later version. Use bmtrain.distributed.all_reduce instead.", DeprecationWarning) - return distributed.all_reduce(loss, "avg") + return distributed.all_reduce(loss, "sum") / config['world_size'] def gather_result(result: torch.Tensor): warnings.warn("bmtrain.gather_result is deprecated and will be removed in later version. Use bmtrain.distributed.all_gather instead.", DeprecationWarning) diff --git a/csrc/cuda/adam.cu b/csrc/cuda/adam.cu index 4b3df9c7..c7bff5ca 100644 --- a/csrc/cuda/adam.cu +++ b/csrc/cuda/adam.cu @@ -26,7 +26,7 @@ __global__ void adam_fp32_accum( float local_m = beta1 * __half2float(m[col]) + (1 - beta1) * local_g; // real_m * scale float local_v = beta2 * v[col] + (1 - beta2) * local_g * local_g / scale; // real_v * scale float local_p = param[col]; - local_p = local_p - lr * local_m / bias_correction1 / (sqrtf(local_v * scale / bias_correction2) + eps) - lr * weight_decay * local_p; + local_p = local_p - lr * local_m / bias_correction1 / (sqrtf(local_v * scale / bias_correction2) + eps * scale) - lr * weight_decay * local_p; param_h[col] = __float2half(local_p); param[col] = local_p; diff --git a/example/train.py b/example/train.py index 0f0454bb..7bc92400 100644 --- a/example/train.py +++ b/example/train.py @@ -54,8 +54,8 @@ def main(): optimizer = bmt.optim.AdamOffloadOptimizer(model.parameters(), weight_decay=1e-2) lr_scheduler = bmt.lr_scheduler.Noam(optimizer, start_lr=1e-3, warmup_iter=40, end_iter=1000, num_iter=0) - loss_scaler = bmt.optim.LossScaler(loss_scale=2**20) - loss_scaler.add_optimizer(optimizer, lr_scheduler) + optim_manager = bmt.optim.OptimManager(loss_scale=2**20) + optim_manager.add_optimizer(optimizer, lr_scheduler) bmt.synchronize() @@ -65,7 +65,6 @@ def main(): for iteration in range(1000): # load data st = time.time() - loss_scaler.zero_grad() with bmt.inspect.inspect_tensor() as inspector: pos = torch.arange(enc_input.size(1)).long().cuda().repeat(enc_input.size(0), 1) @@ -80,8 +79,9 @@ def main(): global_loss = bmt.sum_loss(loss).item() - loss = loss_scaler(loss) - loss.backward() + optim_manager.zero_grad() + + optim_manager.backward(loss) # print inspected tensors in the forward & backward pass # print parameters of the model @@ -96,9 +96,8 @@ def main(): bmt.inspect.inspect_model(model, "*") ) ) - - loss_scaler.step() + optim_manager.step() # record time and loss iteration_time = time.time() - st @@ -113,7 +112,7 @@ def main(): global_loss, avg_loss_recorder.value, lr_scheduler.current_lr, - loss_scaler.loss_scale, + optim_manager.loss_scale, avg_time_recorder.value ) ) diff --git a/tests/test_all.py b/tests/test_all.py new file mode 100644 index 00000000..7dd2c0b0 --- /dev/null +++ b/tests/test_all.py @@ -0,0 +1,29 @@ +import subprocess +from tqdm import tqdm + + +tq = tqdm([ + ("load_ckpt", 1), + ("init_parameters", 1), + ("init_parameters_multi_gpu", 4), + + ("requires_grad", 1), + ("has_inf_nan", 1), + ("dropout", 1), + ("loss_func", 1), + + ("middle_hidden", 4), + ("model_wrapper", 4), + ("send_recv", 4), + ("nccl_backward", 4), + + ("trainging", 4), +]) + +for t, num_gpu in tq: + PREFIX = f"python3 -m torch.distributed.launch --nnodes=1 --node_rank=0 --nproc_per_node={num_gpu} --master_addr=localhost --master_port=32123" + SUFFIX = f"> test_log.txt 2>&1" + command = f"{PREFIX} test_{t}.py {SUFFIX}" + completedProc = subprocess.run(command, shell=True) + assert completedProc.returncode == 0, f"test {t} failed, see test_log.txt for more details." + print(f"finish testing {t}") diff --git a/tests/test_backward.py b/tests/test_backward.py deleted file mode 100644 index 6bda8070..00000000 --- a/tests/test_backward.py +++ /dev/null @@ -1,84 +0,0 @@ -import torch -import bmtrain -from tqdm import tqdm - -from bmtrain.utils import print_rank - -class DistributedLinear(bmtrain.DistributedModule): - def __init__(self, dim_in, dim_out): - super().__init__() - self.dim_in = dim_in - self.dim_out = dim_out - - self.weight = bmtrain.DistributedParameter(torch.empty(dim_in, dim_out, dtype=torch.float)) - - def forward(self, x : torch.Tensor): - last_dim = x.size(-1) - x_viewd = x.view(-1, last_dim) - out = torch.matmul(x_viewd, self.weight).view(x.size()[:-1] + (self.dim_out,)) - var_out = out.var(dim=-1, keepdim=True) - return out / var_out + x - -class TestMLP(bmtrain.DistributedModule): - def __init__(self): - super().__init__() - - self.mlp = torch.nn.ModuleList([ - DistributedLinear(8, 8) for _ in range(16) - ]) - - @bmtrain.checkpoint - def forward_segments(self, layer_ids, hidden_state : torch.Tensor): - for idx in layer_ids: - hidden_state = self.mlp[idx](hidden_state) - return hidden_state - - def forward(self, x): - for segments in [ - [0, 1], [2, 3], - [4, 5], [6, 7], - [8, 9], [10, 11], - [12, 13], [14, 15] - ]: - x = self.forward_segments(segments, x) - bmtrain.wait_loader() - return x - -def main(): - bmtrain.init_distributed() - - model = TestMLP() - - print_rank("Model mem\n", torch.cuda.memory_summary()) - state = model.state_dict() - for kw, param in state.items(): - torch.nn.init.normal_(param, std=0.02) - model.load_state_dict(state) # split parameters automatically - print_rank("Model after loading state dict\n", torch.cuda.memory_summary()) - - for i in range(bmtrain.world_size()): - v = torch.randn(4, 8, dtype=torch.float, device="cuda") - if i == bmtrain.rank(): - raw_data = v - - optimizer = torch.optim.Adam(model.parameters(), lr=1e-4) - for iter in tqdm(range(100)): - # load data - data = raw_data.clone().requires_grad_() - optimizer.zero_grad() - output = model(data) - loss = ((output - bmtrain.rank()) ** 2).mean() - print_rank("Iter %d, loss: " % iter, bmtrain.sum_loss(loss).item()) - - loss.backward() - optimizer.step() - bmtrain.wait_optimizer() # loader stream wait for optimizer - - print_rank(torch.cuda.memory_summary()) - bmtrain.synchronize() - print_rank(torch.cuda.memory_summary(), rank=1) - - - -if __name__ == "__main__": - main() \ No newline at end of file diff --git a/tests/test_dropout.py b/tests/test_dropout.py index d72ee75f..f02da4e3 100644 --- a/tests/test_dropout.py +++ b/tests/test_dropout.py @@ -1,3 +1,5 @@ +from utils import * + import torch import bmtrain as bmt @@ -8,7 +10,6 @@ def __init__(self): self.drop = torch.nn.Dropout(p=0.5) def forward(self, x): - bmt.print_rank(x) return self.drop(x) class OutterModule(bmt.DistributedModule): @@ -21,17 +22,24 @@ def __init__(self) -> None: ]) def forward(self, x): - bmt.print_rank(self.training) return self.blk(x) -def main(): - bmt.init_distributed() - +def test_main(): model = OutterModule() - x = torch.ones(32, device="cuda") - y = model(x) - bmt.print_rank(y) + for _ in range(5): + model.train() + x = torch.ones(32, device="cuda") + y = model(x) + print(y) + assert_neq(x.numel()-y.nonzero().size(0), 0) + + model.eval() + x = torch.ones(32, device="cuda") + y = model(x) + print(y) + assert_eq(x.numel()-y.nonzero().size(0), 0) if __name__ == "__main__": - main() \ No newline at end of file + bmt.init_distributed() + test_main() \ No newline at end of file diff --git a/tests/test_hasinfnan.py b/tests/test_has_inf_nan.py similarity index 93% rename from tests/test_hasinfnan.py rename to tests/test_has_inf_nan.py index b310e546..b1b9b4a9 100644 --- a/tests/test_hasinfnan.py +++ b/tests/test_has_inf_nan.py @@ -1,3 +1,5 @@ +from utils import * + import torch import bmtrain.optim._cuda as G import random @@ -5,7 +7,7 @@ def check(x, v): out = torch.zeros(1, dtype=torch.uint8, device="cuda")[0] G.f_has_inf_nan(x, out) - assert(out.item() == v) + assert_eq(out.item(), v) def test_main(): for i in list(range(1, 100)) + [1000]*10 + [10000]*10 + [100000]*10 + [1000000]*10: diff --git a/tests/test_init_parameters.py b/tests/test_init_parameters.py new file mode 100644 index 00000000..c23249ff --- /dev/null +++ b/tests/test_init_parameters.py @@ -0,0 +1,219 @@ +from utils import * + +import torch +import torch.nn.functional as F +import bmtrain as bmt + +def manual_seed(seed=33): + torch.manual_seed(seed) + import random as random + random.seed(seed) + try: + import numpy as np + np.random.seed(seed) + except ModuleNotFoundError: + pass + +class Linear_NormalInitBefore(torch.nn.Module): + def __init__(self, in_features : int, out_features: int, bias: bool = True, dtype = None) -> None: + super().__init__() + + self.in_features = in_features + self.out_features = out_features + w = torch.empty(out_features, in_features, dtype=dtype, device="cuda") # use cuda to match random algorithm + torch.nn.init.xavier_normal_(w) + self.weight = torch.nn.Parameter(w) + if bias: + b = torch.empty(out_features, dtype=dtype, device="cuda") # use cuda to match random algorithm + torch.nn.init.zeros_(b) + self.bias = torch.nn.Parameter(b) + else: + self.register_parameter('bias', None) + + def forward(self, input): + return F.linear(input, self.weight, self.bias) + +class Linear_NormalInitAfter(torch.nn.Module): + def __init__(self, in_features : int, out_features: int, bias: bool = True, dtype = None) -> None: + super().__init__() + + self.in_features = in_features + self.out_features = out_features + self.weight = torch.nn.Parameter(torch.empty(out_features, in_features, dtype=dtype, device="cuda")) # use cuda to match random algorithm + torch.nn.init.xavier_normal_(self.weight) + if bias: + self.bias = torch.nn.Parameter(torch.empty(out_features, dtype=dtype, device="cuda")) # use cuda to match random algorithm + torch.nn.init.zeros_(self.bias) + else: + self.register_parameter('bias', None) + + def forward(self, input): + return F.linear(input, self.weight, self.bias) + +class Linear_BMTInitializer(bmt.DistributedModule): + def __init__(self, in_features : int, out_features: int, bias: bool = True, dtype = None) -> None: + super().__init__() + + self.in_features = in_features + self.out_features = out_features + self.weight = bmt.DistributedParameter(torch.empty(out_features, in_features, dtype=dtype), init_method=torch.nn.init.xavier_normal_) + if bias: + self.bias = bmt.DistributedParameter(torch.empty(out_features, dtype=dtype), init_method=torch.nn.init.zeros_) + else: + self.register_parameter('bias', None) + + def forward(self, input): + return F.linear(input, self.weight, self.bias) + +class Linear_ManualInitBefore(bmt.DistributedModule): + def __init__(self, in_features : int, out_features: int, bias: bool = True, dtype = None) -> None: + super().__init__() + + self.in_features = in_features + self.out_features = out_features + w = torch.empty(out_features, in_features, dtype=dtype, device="cuda") # use cuda to match random algorithm + torch.nn.init.xavier_normal_(w) + self.weight = bmt.DistributedParameter(w) + if bias: + b = torch.empty(out_features, dtype=dtype, device="cuda") # use cuda to match random algorithm + torch.nn.init.zeros_(b) + self.bias = bmt.DistributedParameter(b) + else: + self.register_parameter('bias', None) + + def forward(self, input): + return F.linear(input, self.weight, self.bias) + +class Linear_ManualInitAfter(bmt.DistributedModule): + def __init__(self, in_features : int, out_features: int, bias: bool = True, dtype = None) -> None: + super().__init__() + + self.in_features = in_features + self.out_features = out_features + self.weight = bmt.DistributedParameter(torch.empty(out_features, in_features, dtype=dtype, device="cuda")) # use cuda to match random algorithm + # torch.nn.init.xavier_normal_(self.weight) + if bias: + self.bias = bmt.DistributedParameter(torch.empty(out_features, dtype=dtype, device="cuda")) # use cuda to match random algorithm + # torch.nn.init.zeros_(self.bias) + else: + self.register_parameter('bias', None) + + def forward(self, input): + return F.linear(input, self.weight, self.bias) + + def extra_repr(self) -> str: + return 'in_features={}, out_features={}, bias={}'.format( + self.in_features, self.out_features, self.bias is not None + ) + +class Linear_Pipeline(bmt.DistributedModule): + def __init__(self, in_features : int, out_features: int, bias: bool = True, dtype = None) -> None: + super().__init__() + + self.l = bmt.PipelineTransformerBlockList([ + Linear_BMTInitializer(in_features, out_features, bias, dtype), + Linear_BMTInitializer(in_features, out_features, bias, dtype), + ]) + + def forward(self, input): + return self.l(input) + +class Linear_BlockList(bmt.DistributedModule): + def __init__(self, in_features : int, out_features: int, bias: bool = True, dtype = None) -> None: + super().__init__() + + self.l = bmt.TransformerBlockList([ + Linear_BMTInitializer(in_features, out_features, bias, dtype), + Linear_BMTInitializer(in_features, out_features, bias, dtype), + ]) + + def forward(self, input): + return self.l(input) + +class Linear_CheckpointList(bmt.DistributedModule): + def __init__(self, in_features : int, out_features: int, bias: bool = True, dtype = None) -> None: + super().__init__() + + self.l = torch.nn.ModuleList([ + bmt.CheckpointBlock(Linear_BMTInitializer(in_features, out_features, bias, dtype)), + bmt.CheckpointBlock(Linear_BMTInitializer(in_features, out_features, bias, dtype)), + ]) + + def forward(self, input): + return self.l(input) + +class Linear_Checkpoint(bmt.DistributedModule): + def __init__(self, in_features : int, out_features: int, bias: bool = True, dtype = None) -> None: + super().__init__() + + self.l = bmt.CheckpointBlock(Linear_BMTInitializer(in_features, out_features, bias, dtype)) + + def forward(self, input): + return self.l(input) + +def test_main(): + shape = [3, 5] + # torch + m = [None] * 10 + ret = [None] * 10 + manual_seed(33) + m[0] = Linear_NormalInitBefore(*shape) + ret[0] = (m[0].weight.data, m[0].bias.data) + + manual_seed(33) + m[1] = Linear_NormalInitAfter(*shape) + ret[1] = (m[1].weight.data, m[1].bias.data) + + # bmtrain + manual_seed(33) + m[2] = Linear_BMTInitializer(*shape) + bmt.init_parameters(m[2]) + ret[2] = (m[2].weight.data, m[2].bias.data) + + manual_seed(33) + m[3] = Linear_ManualInitBefore(*shape) + ret[3] = (m[3].weight.data, m[3].bias.data) + + # manual_seed(33) + # mw = Linear_ManualInitAfter(*shape) # not supported + # print(mw.weight.data, mw.bias.data) + + manual_seed(33) + m[4] = bmt.BMTrainModelWrapper(m[0]) + ret[4] = (m[4].weight.data, m[4].bias.data) + + manual_seed(33) + m[5] = bmt.BMTrainModelWrapper(m[1]) + ret[5] = (m[5].weight.data, m[5].bias.data) + + manual_seed(33) + m[6] = Linear_Pipeline(*shape) + bmt.init_parameters(m[6]) + ret[6] = (m[6].l[0].weight.data, m[6].l[0].bias.data) + + manual_seed(33) + m[7] = Linear_BlockList(*shape) + bmt.init_parameters(m[7]) + ret[7] = (m[7].l[0].weight.data, m[7].l[0].bias.data) + + manual_seed(33) + m[8] = Linear_CheckpointList(*shape) + bmt.init_parameters(m[8]) + ret[8] = (m[8].l[0].weight.data, m[8].l[0].bias.data) + + manual_seed(33) + m[9] = Linear_Checkpoint(*shape) + bmt.init_parameters(m[9]) + ret[9] = (m[9].l.weight.data, m[9].l.bias.data) + + for i in range(10): + ret[i] = ( ret[i][0].view(-1), ret[i][1].view(-1) ) + for i in range(10): + for j in range(10): + assert_all_eq(ret[i][0], ret[j][0]) + assert_all_eq(ret[i][1], ret[j][1]) + +if __name__ == "__main__": + bmt.init_distributed(pipe_size=1) + + test_main() \ No newline at end of file diff --git a/tests/test_init_parameters_multi_gpu.py b/tests/test_init_parameters_multi_gpu.py new file mode 100644 index 00000000..1e61568c --- /dev/null +++ b/tests/test_init_parameters_multi_gpu.py @@ -0,0 +1,146 @@ +from utils import * + +import os +import torch +import torch.nn.functional as F +import bmtrain as bmt + +def manual_seed(seed=33): + torch.manual_seed(seed) + import random as random + random.seed(seed) + try: + import numpy as np + np.random.seed(seed) + except ModuleNotFoundError: + pass + +class Linear_Normal(torch.nn.Module): + def __init__(self, in_features : int, out_features: int, bias: bool = True, dtype = None) -> None: + super().__init__() + + self.in_features = in_features + self.out_features = out_features + self.weight = torch.nn.Parameter(torch.empty(out_features, in_features, dtype=dtype, device="cuda")) # use cuda to match random algorithm + torch.nn.init.xavier_normal_(self.weight) + if bias: + self.bias = torch.nn.Parameter(torch.empty(out_features, dtype=dtype, device="cuda")) # use cuda to match random algorithm + torch.nn.init.zeros_(self.bias) + else: + self.register_parameter('bias', None) + + def forward(self, input): + return F.linear(input, self.weight, self.bias) + +class Linear_BMTInitializer(bmt.DistributedModule): + def __init__(self, in_features : int, out_features: int, bias: bool = True, dtype = None) -> None: + super().__init__() + + self.in_features = in_features + self.out_features = out_features + self.weight = bmt.DistributedParameter(torch.empty(out_features, in_features, dtype=dtype), init_method=torch.nn.init.xavier_normal_) + if bias: + self.bias = bmt.DistributedParameter(torch.empty(out_features, dtype=dtype), init_method=torch.nn.init.zeros_) + else: + self.register_parameter('bias', None) + + def forward(self, input): + return F.linear(input, self.weight, self.bias) + +class Linear_NormalList(torch.nn.Module): + def __init__(self, in_features : int, out_features: int, bias: bool = True, dtype = None) -> None: + super().__init__() + + self.l = torch.nn.ModuleList([ + Linear_Normal(in_features, out_features, bias, dtype), + Linear_Normal(in_features, out_features, bias, dtype), + ]) + + def forward(self, input): + return self.l(input) + +class Linear_Pipeline(bmt.DistributedModule): + def __init__(self, in_features : int, out_features: int, bias: bool = True, dtype = None) -> None: + super().__init__() + + self.l = bmt.PipelineTransformerBlockList([ + Linear_BMTInitializer(in_features, out_features, bias, dtype), + Linear_BMTInitializer(in_features, out_features, bias, dtype), + ]) + + def forward(self, input): + return self.l(input) + +class Linear_BlockList(bmt.DistributedModule): + def __init__(self, in_features : int, out_features: int, bias: bool = True, dtype = None) -> None: + super().__init__() + + self.l = bmt.TransformerBlockList([ + Linear_BMTInitializer(in_features, out_features, bias, dtype), + Linear_BMTInitializer(in_features, out_features, bias, dtype), + ]) + + def forward(self, input): + return self.l(input) + +class Linear_CheckpointList(bmt.DistributedModule): + def __init__(self, in_features : int, out_features: int, bias: bool = True, dtype = None) -> None: + super().__init__() + + self.l = torch.nn.ModuleList([ + bmt.CheckpointBlock(Linear_BMTInitializer(in_features, out_features, bias, dtype)), + bmt.CheckpointBlock(Linear_BMTInitializer(in_features, out_features, bias, dtype)), + ]) + + def forward(self, input): + return self.l(input) + +def check(ckpt_path, ckpt_path_ref): + if bmt.rank() == 0: + ckpt1 = torch.load(ckpt_path) + ckpt2 = torch.load(ckpt_path_ref) + for (k1, v1), (k2, v2) in zip(ckpt1.items(), ckpt2.items()): + assert_eq(k1, k2) + print(v1, v2) + assert_all_eq(v1.cuda(), v2.cuda()) + +def test_main(): + ckpt_path_ref = "test_ckpt_ref.pt" + ckpt_path = "test_ckpt.pt" + shape = [3, 5] + # torch + m = [None] * 4 + ret = [None] * 4 + + manual_seed(33) + m[0] = Linear_NormalList(*shape) + if bmt.rank() == 0: + torch.save(m[0].state_dict(), ckpt_path_ref) + + manual_seed(33) + m[1] = Linear_Pipeline(*shape) + bmt.init_parameters(m[1]) + bmt.save(m[1], ckpt_path) + check(ckpt_path, ckpt_path_ref) + + # bmtrain + manual_seed(33) + m[2] = Linear_BlockList(*shape) + bmt.init_parameters(m[2]) + bmt.save(m[2], ckpt_path) + check(ckpt_path, ckpt_path_ref) + + manual_seed(33) + m[3] = Linear_CheckpointList(*shape) + bmt.init_parameters(m[3]) + bmt.save(m[3], ckpt_path) + check(ckpt_path, ckpt_path_ref) + + if bmt.rank() == 0: + os.remove(ckpt_path) + os.remove(ckpt_path_ref) + +if __name__ == "__main__": + bmt.init_distributed(pipe_size=2) + + test_main() \ No newline at end of file diff --git a/tests/test_load_ckpt.py b/tests/test_load_ckpt.py new file mode 100644 index 00000000..56298940 --- /dev/null +++ b/tests/test_load_ckpt.py @@ -0,0 +1,66 @@ +from utils import * +import torch +import torch.nn.functional as F +import bmtrain as bmt +import os + +class Linear_Normal(torch.nn.Module): + def __init__(self, in_features : int, out_features: int, bias: bool = True, dtype = None) -> None: + super().__init__() + + self.in_features = in_features + self.out_features = out_features + self.weight = torch.nn.Parameter(torch.empty(out_features, in_features, dtype=dtype)) + torch.nn.init.xavier_normal_(self.weight) + if bias: + self.bias = torch.nn.Parameter(torch.empty(out_features, dtype=dtype)) + torch.nn.init.zeros_(self.bias) + else: + self.register_parameter('bias', None) + + def forward(self, input): + return F.linear(input, self.weight, self.bias) + +class Linear_BMT(bmt.DistributedModule): + def __init__(self, in_features : int, out_features: int, bias: bool = True, dtype = None) -> None: + super().__init__() + + self.in_features = in_features + self.out_features = out_features + self.weight = bmt.DistributedParameter(torch.empty(out_features, in_features, dtype=dtype), init_method=torch.nn.init.xavier_normal_) + if bias: + self.bias = bmt.DistributedParameter(torch.empty(out_features, dtype=dtype), init_method=torch.nn.init.zeros_) + else: + self.register_parameter('bias', None) + + def forward(self, input): + return F.linear(input, self.weight, self.bias) + +def test_main(): + ckpt_path = "test_ckpt.pt" + m = Linear_Normal(7, 5).cuda() + m2 = Linear_BMT(7, 5) + torch.save(m.state_dict(), ckpt_path) + bmt.load(m2, ckpt_path) + + print(m.weight.data) + print(m2.weight.data) + assert_all_eq(m.weight.data, m2.weight.data) + assert_all_eq(m.bias.data, m2.bias.data) + + os.remove(ckpt_path) + + bmt.save(m2, ckpt_path) + m.load_state_dict(torch.load(ckpt_path)) + + print(m.weight.data) + print(m2.weight.data) + assert_all_eq(m.weight.data, m2.weight.data) + assert_all_eq(m.bias.data, m2.bias.data) + + os.remove(ckpt_path) + +if __name__ == "__main__": + bmt.init_distributed() + + test_main() \ No newline at end of file diff --git a/tests/test_loss.py b/tests/test_loss_func.py similarity index 74% rename from tests/test_loss.py rename to tests/test_loss_func.py index ac7692de..a76be5f4 100644 --- a/tests/test_loss.py +++ b/tests/test_loss_func.py @@ -1,8 +1,9 @@ +from utils import * + import torch import bmtrain as bmt import torch import random -import time import copy def run(x, tgt, loss_func, bigmodel=None, scale=32768, use_float=False): @@ -13,28 +14,23 @@ def run(x, tgt, loss_func, bigmodel=None, scale=32768, use_float=False): if bigmodel is not None: bigmodel = bigmodel.float() x = x.requires_grad_() - # print(torch.cuda.memory_summary()) if bigmodel is None: loss = loss_func(x, tgt) else: t = bigmodel(x) loss = loss_func(t, tgt) - # loss = loss_func(bigmodel(x), tgt) - # print(torch.cuda.memory_summary()) (loss * scale).backward() - # print(torch.cuda.memory_summary()) return loss, x.grad def check(x, tgt, loss_func1, loss_func2, bigmodel=None): loss_1, grad_1 = run(x, tgt, loss_func1, bigmodel=bigmodel) loss_2, grad_2 = run(x, tgt, loss_func2, bigmodel=bigmodel, use_float=True) - assert(grad_1.isnan().sum()==0) - assert(grad_2.isnan().sum()==0) - print((loss_1 - loss_2).abs().item()) - print((grad_1 - grad_2).abs().max().item()) - print() + assert_eq(grad_1.isnan().sum(), 0) + assert_eq(grad_2.isnan().sum(), 0) + assert_lt((loss_1 - loss_2).abs().item(), 1e-5) + assert_lt((grad_1 - grad_2).abs().max().item(), 1e-2) -def simple_test(): +def test_simple(): loss_func1 = bmt.loss.FusedCrossEntropy() loss_func2 = torch.nn.CrossEntropyLoss() @@ -83,24 +79,7 @@ def test_inplace(): tgt = torch.randint(0, C, (N,)).cuda().long() check(x, tgt, loss_func1, loss_func2, bigmodel=bigmodel) - -def benchmark_memory(): - N = 32 * 512 - C = 50000 - - bigmodel = torch.nn.Linear(1, C).cuda().half() - x = torch.randn(N, 1).cuda().half() - tgt = torch.randint(0, C, (N,)).cuda().long() - - # loss_func = bmt.loss.FusedCrossEntropy(inplace=True) - loss_func = bmt.loss.FusedCrossEntropy(inplace=False) - # loss_func = torch.nn.CrossEntropyLoss() - - loss, grad = run(x, tgt, loss_func, bigmodel=bigmodel) - print(loss, grad) - if __name__ == "__main__": test_other() test_inplace() - simple_test() - # benchmark_memory() \ No newline at end of file + test_simple() \ No newline at end of file diff --git a/tests/test_middle_hidden.py b/tests/test_middle_hidden.py index ca0dc9b3..2073588d 100644 --- a/tests/test_middle_hidden.py +++ b/tests/test_middle_hidden.py @@ -1,3 +1,5 @@ +from utils import * + import bmtrain as bmt import random import torch @@ -115,7 +117,7 @@ def manual_seed(seed=33): except ModuleNotFoundError: pass -def sub_test(name, cls, num_layer, dim, batch, seq_len, only_last=False, only_middle=False, mix_test=False): +def sub_run(name, cls, num_layer, dim, batch, seq_len, only_last=False, only_middle=False, mix_test=False): manual_seed() pre = Linear(dim, dim) @@ -135,15 +137,14 @@ def sub_test(name, cls, num_layer, dim, batch, seq_len, only_last=False, only_mi bmt.init_parameters(m) m = cls(pre, [m for m in ms], post) + ret = "" if only_last: logits = m(inp) loss = (logits * last_weight).sum() loss.backward() - bmt.print_rank(f"========================{name}:only last========================") - bmt.print_rank( - bmt.inspect.format_summary( - bmt.inspect.inspect_model(m, '*') - ) + ret += f"========================only last========================\n" + ret += bmt.inspect.format_summary( + bmt.inspect.inspect_model(m, '*') ) if only_middle: logits, hidden_states = m(inp, return_hidden_states=True) @@ -152,11 +153,9 @@ def sub_test(name, cls, num_layer, dim, batch, seq_len, only_last=False, only_mi for i, hidden_state in enumerate(hidden_states) ]) loss.backward() - bmt.print_rank(f"========================{name}:only middle========================") - bmt.print_rank( - bmt.inspect.format_summary( - bmt.inspect.inspect_model(m, '*') - ) + ret += f"========================only middle========================\n" + ret += bmt.inspect.format_summary( + bmt.inspect.inspect_model(m, '*') ) if mix_test: logits, hidden_states = m(inp, return_hidden_states=True) @@ -165,24 +164,33 @@ def sub_test(name, cls, num_layer, dim, batch, seq_len, only_last=False, only_mi for i, hidden_state in enumerate(hidden_states) ]) + (logits * last_weight).sum() loss.backward() - bmt.print_rank(f"========================{name}:mix========================") - bmt.print_rank( - bmt.inspect.format_summary( - bmt.inspect.inspect_model(m, '*') - ) + ret += f"========================mix========================\n" + ret += bmt.inspect.format_summary( + bmt.inspect.inspect_model(m, '*') ) + return ret.replace("None ", "0.0000") + "\n" # replace for matching None grad with zero_grad -def test(name, cls, num_layer=4, dim=4096, batch=32, seq_len=256): - sub_test(name, cls, num_layer=num_layer, dim=dim, batch=batch, seq_len=seq_len, only_last=True) +def run(name, cls, num_layer=4, dim=4096, batch=32, seq_len=256): + ret = "" + ret += sub_run(name, cls, num_layer=num_layer, dim=dim, batch=batch, seq_len=seq_len, only_last=True) bmt.synchronize() - sub_test(name, cls, num_layer=num_layer, dim=dim, batch=batch, seq_len=seq_len, only_middle=True) + ret += sub_run(name, cls, num_layer=num_layer, dim=dim, batch=batch, seq_len=seq_len, only_middle=True) bmt.synchronize() - sub_test(name, cls, num_layer=num_layer, dim=dim, batch=batch, seq_len=seq_len, mix_test=True) + ret += sub_run(name, cls, num_layer=num_layer, dim=dim, batch=batch, seq_len=seq_len, mix_test=True) bmt.synchronize() - -bmt.init_distributed(pipe_size=4) - -test("normal", Model_NORMAL) -test("block", Model_BLOCK) -test("zero", Model_ZERO) -test("pipe", Model_PIPE) \ No newline at end of file + return ret + +def test_main(): + ret = [] + ret.append( run("normal", Model_NORMAL) ) + ret.append( run("block", Model_BLOCK) ) + ret.append( run("zero", Model_ZERO) ) + ret.append( run("pipe", Model_PIPE) ) + for r in ret: + bmt.prnit_rank(r) + for r in ret: + for r2 in ret: + assert_eq(r, r2) + +if __name__ == "__main__": + bmt.init_distributed(pipe_size=4) diff --git a/tests/test_wrapper.py b/tests/test_model_wrapper.py similarity index 72% rename from tests/test_wrapper.py rename to tests/test_model_wrapper.py index e2b0ea11..409107e3 100644 --- a/tests/test_wrapper.py +++ b/tests/test_model_wrapper.py @@ -1,3 +1,5 @@ +from utils import * + from typing import Optional import torch import math @@ -169,11 +171,7 @@ def forward(self, return logits -def main(): - bmt.init_distributed( - seed=0, - ) - +def test_main(): model = GPT( num_layers=8, vocab_size=10240, @@ -186,16 +184,10 @@ def main(): dtype=torch.half ) - model.load_state_dict(torch.load("../example/ckpt-0.pt")) - - model = bmt.BMTrainModelWrapper(model) + bmt_model = bmt.BMTrainModelWrapper(model) + model = model.cuda() - bmt.print_rank("Model memory") - bmt.print_rank(torch.cuda.memory_summary()) - bmt.synchronize() - - # data - # generate dummy data for each rank + # use the break if i == bmt.rank() to generate different data on different rank torch.manual_seed(1234) batch_size = 2 seq_len = 512 @@ -214,80 +206,16 @@ def main(): if i == bmt.rank(): break - - loss_func = torch.nn.CrossEntropyLoss(ignore_index=-100) - optimizer = bmt.optim.AdamOffloadOptimizer(model.parameters(), weight_decay=1e-2) - lr_scheduler = bmt.lr_scheduler.Noam(optimizer, start_lr=1e-3, warmup_iter=40, end_iter=1000, num_iter=0) - - loss_scaler = bmt.optim.LossScaler(loss_scale=2**20) - loss_scaler.add_optimizer(optimizer, lr_scheduler) + pos = torch.arange(enc_input.size(1)).long().cuda().repeat(enc_input.size(0), 1) + logits = model(enc_input, pos, pos < enc_length[:, None]) + bmt_logits = bmt_model(enc_input, pos, pos < enc_length[:, None]) + print(logits) bmt.synchronize() - - avg_time_recorder = bmt.utils.AverageRecorder() - avg_loss_recorder = bmt.utils.AverageRecorder() - - for iteration in range(1000): - # load data - st = time.time() - loss_scaler.zero_grad() - - with bmt.inspect.inspect_tensor() as inspector: - pos = torch.arange(enc_input.size(1)).long().cuda().repeat(enc_input.size(0), 1) - logits = model( - enc_input, - pos, - pos < enc_length[:, None] - ) - batch, seq_len, vocab_out_size = logits.size() - - loss = loss_func(logits.view(batch * seq_len, vocab_out_size), targets.view(batch * seq_len)) - - global_loss = bmt.sum_loss(loss).item() - - loss = loss_scaler(loss) - loss.backward() - - # print inspected tensors in the forward & backward pass - # print parameters of the model - if iteration % 100 == 0: - bmt.print_rank( - bmt.inspect.format_summary( - inspector.get_summary() - ) - ) - bmt.print_rank( - bmt.inspect.format_summary( - bmt.inspect.inspect_model(model, "*") - ) - ) - - - loss_scaler.step() - - # record time and loss - iteration_time = time.time() - st - - avg_time_recorder.record(iteration_time) - avg_loss_recorder.record(global_loss) - - # print time and loss - bmt.print_rank( - "| Iter: {:6d} | loss: {:.4f} average_loss: {:.4f} | lr: {:.4e} scale: {:10.4f} | time: {:.4f}".format( - iteration, - global_loss, - avg_loss_recorder.value, - lr_scheduler.current_lr, - loss_scaler.loss_scale, - avg_time_recorder.value - ) - ) - - # save model - if iteration % 1000 == 0: - bmt.save(model, "ckpt-%d.pt" % iteration) - - bmt.save(model, "checkpoint.pt") + print(bmt_logits) + assert_all_eq(logits, bmt_logits) if __name__ == '__main__': - main() \ No newline at end of file + bmt.init_distributed(seed=0) + + test_main() \ No newline at end of file diff --git a/tests/test_dist.py b/tests/test_nccl_backward.py similarity index 64% rename from tests/test_dist.py rename to tests/test_nccl_backward.py index c6844ee3..3dcd0560 100644 --- a/tests/test_dist.py +++ b/tests/test_nccl_backward.py @@ -1,15 +1,20 @@ +from utils import * + import bmtrain as bmt import torch -def main(): - bmt.init_distributed() +def test_main(): x = torch.full((1,), bmt.rank() + 1, dtype=torch.half, device="cuda").requires_grad_(True) y = bmt.distributed.all_reduce(x, "prod").view(-1) - bmt.print_rank(y) loss = (y * y).sum() / 2 loss.backward() + ref = y + for i in range(bmt.world_size()): + if i != bmt.rank(): ref *= i+1 print(x.grad) - + assert_eq(x.grad, ref) if __name__ == "__main__": - main() \ No newline at end of file + bmt.init_distributed() + + test_main() \ No newline at end of file diff --git a/tests/test_requires_grad.py b/tests/test_requires_grad.py index 2a9bab43..707bd793 100644 --- a/tests/test_requires_grad.py +++ b/tests/test_requires_grad.py @@ -1,3 +1,5 @@ +from utils import * + import bmtrain as bmt import torch from bmtrain import config @@ -32,27 +34,40 @@ def run(m, a, b): loss = logits.sum() loss.backward() - bmt.print_rank( - bmt.inspect.format_summary( + sm = bmt.inspect.format_summary( bmt.inspect.inspect_model(m, '*') ) - ) - print(a.weight.grad is None) - print(a.bias.grad is None) - -bmt.init_distributed() -a = Linear(256, 256) -b = Linear(256, 256) -m = TransformerBlockList([CheckpointBlock(a), CheckpointBlock(b)]) -bmt.init_parameters(m) - -a.bias.requires_grad_(False) -run(m, a, b) - -a.weight.requires_grad_(False) -a.bias.requires_grad_(True) -run(m, a, b) - -a.weight.requires_grad_(True) -a.bias.requires_grad_(False) -run(m, a, b) \ No newline at end of file + return a.weight.grad is None, a.bias.grad is None, sm + +def test_main(): + a = Linear(256, 256) + b = Linear(256, 256) + m = TransformerBlockList([CheckpointBlock(a), CheckpointBlock(b)]) + bmt.init_parameters(m) + + a.bias.requires_grad_(False) + awg, abg, sm1 = run(m, a, b) + print(awg, abg, sm1) + assert_eq((awg, abg), (False, True)) + assert_eq(sm1.split('\n')[2].split()[-2:], ["0.0000", "0.0000"]) + + a.weight.requires_grad_(False) + a.bias.requires_grad_(True) + awg, abg, sm2 = run(m, a, b) + print(awg, abg, sm2) + assert_eq((awg, abg), (False, False)) + assert_eq(sm1.split('\n')[1], sm2.split('\n')[1]) + assert_neq(sm1.split('\n')[2], sm2.split('\n')[2]) + + a.weight.requires_grad_(True) + a.bias.requires_grad_(False) + awg, abg, sm3 = run(m, a, b) + print(awg, abg, sm3) + assert_eq((awg, abg), (False, False)) + assert_neq(sm2.split('\n')[1], sm3.split('\n')[1]) + assert_eq(sm2.split('\n')[2], sm3.split('\n')[2]) + +if __name__ == "__main__": + bmt.init_distributed() + + test_main() \ No newline at end of file diff --git a/tests/test_send_recv.py b/tests/test_send_recv.py index f969d661..95c9c1e5 100644 --- a/tests/test_send_recv.py +++ b/tests/test_send_recv.py @@ -1,27 +1,22 @@ -import bmtrain as bmt +from utils import * + import torch +import bmtrain as bmt from bmtrain.global_var import config -from bmtrain.pipe_comm import send_activations, recv_activations, gather_input -from bmtrain import nccl -from torch.distributed.distributed_c10d import recv -from time import sleep -def test_send_tensor(): - bmt.init_distributed() - current_stream = torch.cuda.current_stream() - groups = [0,2] - rank = config['rank'] -def test_gather_input(): - bmt.init_distributed(pipe_size=2) - if config['topology'].get_group_id("pipe") == 0: - a = torch.ones((2,1)) + +def test_send_recv(): + if config["topology"].stage_id == 0: + a = torch.ones((2,1)) * (config["zero_rank"]+1) + a = a.cuda() + print(f"send {a}") + bmt.distributed.send_activations(a, 1, config["pipe_comm"]) else: - a = torch.zeros((2,1)) - res = gather_input(a.cuda(),config['pipe_comm']) - if config['topology'].get_group_rank("pipe") == 0: - print(res) - -def main(): - test_send_tensor() + ref = torch.ones((2,1)) * (config["zero_rank"]+1) + a = bmt.distributed.recv_activations(0, config["pipe_comm"]) + print(f"recv {a}") + assert_all_eq(a, ref.cuda()) if __name__ == '__main__': - main() \ No newline at end of file + bmt.init_distributed(pipe_size=2) + + test_send_recv() \ No newline at end of file diff --git a/tests/test_training.py b/tests/test_training.py new file mode 100644 index 00000000..7342fe6c --- /dev/null +++ b/tests/test_training.py @@ -0,0 +1,445 @@ +from bmtrain.optim import optim_manager +from utils import * + +from typing import Optional +import torch +import math +import torch.nn.functional as F +import bmtrain as bmt +from bmtrain.global_var import config +import os + +class Attention(torch.nn.Module): + def __init__(self, + dim_model : int, dim_head : int, + num_heads : int, bias : bool = True, + dtype = None + ) -> None: + super().__init__() + + self.project_q = torch.nn.Linear(dim_model, dim_head * num_heads, bias=bias, dtype=dtype) + self.project_k = torch.nn.Linear(dim_model, dim_head * num_heads, bias=bias, dtype=dtype) + self.project_v = torch.nn.Linear(dim_model, dim_head * num_heads, bias=bias, dtype=dtype) + + self.project_out = torch.nn.Linear(dim_head * num_heads, dim_model, bias=bias, dtype=dtype) + + self.softmax = torch.nn.Softmax(dim=-1) + self.num_heads = num_heads + self.dim_head = dim_head + self.dim_model = dim_model + + def forward(self, + hidden_q : torch.Tensor, # (batch_size, seq_q, dim_model) + hidden_kv : torch.Tensor, # (batch_size, seq_kv, dim_model) + mask : torch.BoolTensor, # (batch_size, seq_q, seq_kv) + position_bias : Optional[torch.Tensor] = None, # (batch, num_heads, seq_q, seq_kv) + ) -> torch.Tensor: + batch_size, seq_q, dim_model = hidden_q.size() + seq_kv = hidden_kv.size(1) + + h_q : torch.Tensor = self.project_q(hidden_q) + h_k : torch.Tensor = self.project_k(hidden_kv) + h_v : torch.Tensor = self.project_v(hidden_kv) + + h_q = h_q.view(batch_size, seq_q, self.num_heads, self.dim_head) + h_k = h_k.view(batch_size, seq_kv, self.num_heads, self.dim_head) + h_v = h_v.view(batch_size, seq_kv, self.num_heads, self.dim_head) + + h_q = h_q.permute(0, 2, 1, 3).contiguous() + h_k = h_k.permute(0, 2, 1, 3).contiguous() + h_v = h_v.permute(0, 2, 1, 3).contiguous() + + h_q = h_q.view(batch_size * self.num_heads, seq_q, self.dim_head) + h_k = h_k.view(batch_size * self.num_heads, seq_kv, self.dim_head) + h_v = h_v.view(batch_size * self.num_heads, seq_kv, self.dim_head) + + score = torch.bmm( + h_q, h_k.transpose(1, 2) + ) + score = score / math.sqrt(self.dim_head) + + score = score.view(batch_size, self.num_heads, seq_q, seq_kv) + + if position_bias is not None: + score = score + position_bias.view(batch_size, self.num_heads, seq_q, seq_kv) + + score = torch.where( + mask.view(batch_size, 1, seq_q, seq_kv), + score, + torch.scalar_tensor(float('-inf'), device=score.device, dtype=score.dtype) + ) + + score = torch.where( + mask.view(batch_size, 1, seq_q, seq_kv), + self.softmax(score), + torch.scalar_tensor(0, device=score.device, dtype=score.dtype) + ) + + score = score.view(batch_size * self.num_heads, seq_q, seq_kv) + + h_out = torch.bmm( + score, h_v + ) + h_out = h_out.view(batch_size, self.num_heads, seq_q, self.dim_head) + h_out = h_out.permute(0, 2, 1, 3).contiguous() + h_out = h_out.view(batch_size, seq_q, self.num_heads * self.dim_head) + + attn_out = self.project_out(h_out) + return attn_out + +class Feedforward(torch.nn.Module): + def __init__(self, dim_model : int, dim_ff : int, bias : bool = True, dtype = None) -> None: + super().__init__() + + self.w_in = torch.nn.Linear(dim_model, dim_ff, bias = bias, dtype=dtype) + self.w_out = torch.nn.Linear(dim_ff, dim_model, bias = bias, dtype=dtype) + + self.relu = torch.nn.ReLU() + + def forward(self, input : torch.Tensor) -> torch.Tensor: + return self.w_out(self.relu(self.w_in(input))) + + +class TransformerEncoder(torch.nn.Module): + def __init__(self, + dim_model : int, dim_head : int, num_heads : int, dim_ff : int, + bias : bool = True, dtype = None + ) -> None: + super().__init__() + + self.ln_attn = torch.nn.LayerNorm(dim_model, dtype=dtype) + self.attn = Attention(dim_model, dim_head, num_heads, bias=bias, dtype=dtype) + + self.ln_ff = torch.nn.LayerNorm(dim_model, dtype=dtype) + self.ff = Feedforward(dim_model, dim_ff, bias=bias, dtype=dtype) + + def forward(self, + hidden : torch.Tensor, # (batch, seq_len, dim_model) + mask : torch.BoolTensor, # (batch, seq_len, dim_model) + position_bias : Optional[torch.Tensor] = None, # (batch, num_head, seq_len, seq_len) + ): + x = self.ln_attn(hidden) + x = self.attn(x, x, mask, position_bias) + hidden = hidden + x + + x = self.ln_ff(hidden) + x = self.ff(x) + hidden = hidden + x + + return hidden + + +class GPT(torch.nn.Module): + def __init__(self, + num_layers : int, vocab_size : int, + dim_model : int, dim_head : int, num_heads : int, dim_ff : int, + max_distance : int, + bias : bool = True, dtype = None + ) -> None: + super().__init__() + + self.dtype = dtype + self.max_distance = max_distance + + self.word_emb = torch.nn.Embedding(vocab_size, dim_model, dtype=dtype) + self.pos_emb = torch.nn.Embedding(max_distance, dim_model, dtype=dtype) + self.dim_model = dim_model + + self.transformers = torch.nn.ModuleList([ + TransformerEncoder( + dim_model, dim_head, num_heads, dim_ff, bias, dtype + ) + for _ in range(num_layers) + ]) + + self.layernorm = torch.nn.LayerNorm(dim_model, dtype=dtype) + + def forward(self, + input : torch.LongTensor, # (batch, seq_len) + pos : torch.LongTensor, # (batch, seq_len) + mask : torch.BoolTensor, # (batch, seq_len) + ) -> torch.Tensor: + + mask_2d = mask[:, None, :] & mask[:, :, None] # (batch, seq_len, seq_len) + mask_2d = mask_2d & (pos[:, None, :] >= pos[:, :, None]) + + input_emb = self.pos_emb(pos) + self.word_emb(input) + + out = input_emb + if isinstance(self.transformers, torch.nn.ModuleList): + for layer in self.transformers: + out = layer(out, mask_2d, None) + else: + out = self.transformers(out, mask_2d, None) + out = self.layernorm(out) + + logits = F.linear(out, self.word_emb.weight) / math.sqrt(self.dim_model) + + return logits + +def sub_train_torch(model, loss_func_cls, optimizer_cls): + loss_func = loss_func_cls(ignore_index=-100) + optimizer = optimizer_cls(model.parameters(), weight_decay=1e-2) + lr_scheduler = bmt.lr_scheduler.Noam(optimizer, start_lr=1e-3, warmup_iter=40, end_iter=1000, num_iter=0) + + optim_manager = bmt.optim.OptimManager(loss_scale=2**20 if model.dtype == torch.half else None) + optim_manager.add_optimizer(optimizer, lr_scheduler) + + # use the break if i == bmt.rank() to generate different data on different rank + torch.manual_seed(2333) + batch_size = 2 + seq_len = 512 + + sents = [] + enc_lengths = [] + enc_inputs = [] + targetss = [] + masks = [] + for i in range(bmt.world_size()): + sent = torch.randint(0, 10240, (batch_size, seq_len + 1)) + enc_length = torch.randint(128, seq_len, (batch_size,)).long().cuda() + enc_input = sent[:, :-1].long().cuda() + targets = sent[:, 1:].long().cuda() + mask = torch.arange(seq_len).long().cuda()[None, :] < enc_length[:, None] + targets = torch.where( + mask, + targets, + torch.full_like(targets, -100, dtype=torch.long) + ) + + sents.append(sent) + enc_lengths.append(enc_length) + enc_inputs.append(enc_input) + targetss.append(targets) + masks.append(mask) + + sent = torch.cat(sents, dim=0) + enc_length = torch.cat(enc_lengths, dim=0) + enc_input = torch.cat(enc_inputs, dim=0) + targets = torch.cat(targetss, dim=0) + mask = torch.cat(masks, dim=0) + + pos = torch.arange(enc_input.size(1)).long().cuda().repeat(enc_input.size(0), 1) + + logs = [] + for iter in range(100): + optim_manager.zero_grad() + + logits = model(enc_input, pos, pos < enc_length[:, None]) + + batch, seq_len, vocab_out_size = logits.size() + + loss = loss_func(logits.view(batch * seq_len, vocab_out_size), targets.view(batch * seq_len)) + + global_loss = loss.item() + + loss = optim_manager.loss_scale * loss + loss.backward() + + grad_norm = optim_manager.clip_grad_norm(optimizer.param_groups, max_norm=10.0) + + optim_manager.step() + + bmt.print_rank("| Iter: {:6d} | loss: {:.4f} {:.4f} | lr: {:.4e} scale: {:10.4f} | grad_norm: {:.4f} |".format( + iter, + global_loss, + loss, + lr_scheduler.current_lr, + optim_manager.loss_scale, + grad_norm, + )) + logs.append(global_loss) + + summary = bmt.inspect.inspect_model(model, "*") + return logs, summary + +def sub_train(model, loss_func_cls, optimizer_cls): + loss_func = loss_func_cls(ignore_index=-100) + optimizer = optimizer_cls(model.parameters(), weight_decay=1e-2) + lr_scheduler = bmt.lr_scheduler.Noam(optimizer, start_lr=1e-3, warmup_iter=40, end_iter=1000, num_iter=0) + + optim_manager = bmt.optim.OptimManager(loss_scale=2**20 if model.dtype == torch.half else None) + optim_manager.add_optimizer(optimizer, lr_scheduler) + + # use the break if i == bmt.rank() to generate different data on different rank + torch.manual_seed(2333) + batch_size = 2 + seq_len = 512 + + for i in range(bmt.world_size()): + sent = torch.randint(0, 10240, (batch_size, seq_len + 1)) + enc_length = torch.randint(128, seq_len, (batch_size,)).long().cuda() + enc_input = sent[:, :-1].long().cuda() + targets = sent[:, 1:].long().cuda() + mask = torch.arange(seq_len).long().cuda()[None, :] < enc_length[:, None] + targets = torch.where( + mask, + targets, + torch.full_like(targets, -100, dtype=torch.long) + ) + + if i == bmt.rank(): + break + + pos = torch.arange(enc_input.size(1)).long().cuda().repeat(enc_input.size(0), 1) + + logs = [] + for iter in range(100): + optim_manager.zero_grad() + + logits = model(enc_input, pos, pos < enc_length[:, None]) + + batch, seq_len, vocab_out_size = logits.size() + + loss = loss_func(logits.view(batch * seq_len, vocab_out_size), targets.view(batch * seq_len)) + + global_loss = bmt.sum_loss(loss).item() + + optim_manager.backward(loss) + + grad_norm = optim_manager.clip_grad_norm(optimizer.param_groups, max_norm=10.0) + + optim_manager.step() + + bmt.print_rank("| Iter: {:6d} | loss: {:.4f} {:.4f} | lr: {:.4e} scale: {:10.4f} | grad_norm: {:.4f} |".format( + iter, + global_loss, + loss, + lr_scheduler.current_lr, + optim_manager.loss_scale, + grad_norm, + )) + logs.append(global_loss) + + summary = bmt.inspect.inspect_model(model, "*") + return logs, summary + +def train(model, loss_func, optimizer): + key = f"{model[0]}*{loss_func[0]}*{optimizer[0]}" + model = model[1]() + if key.startswith("torch"): + ret = sub_train_torch(model, loss_func[1], optimizer[1]) + else: + ret = sub_train(model, loss_func[1], optimizer[1]) + del model + bmt.print_rank(f"finished {key}") + return key, ret + +def test_main(test_fp16=True, test_fp32=True): + ckpt_path = "test_ckpt.pt" + + kwargs = { + "num_layers": 8, + "vocab_size": 10240, + "dim_model": 2560, + "dim_head": 80, + "num_heads": 32, + "dim_ff": 8192, + "max_distance": 1024, + "bias": True, + "dtype": None, + } + + def make_ref_ckpt(): + model = GPT(**kwargs) + if bmt.rank() == 0: + torch.save(model.state_dict(), ckpt_path) + bmt.synchronize() + del model + + ret = {} + def torch_model(): + model = GPT(**kwargs) + model.load_state_dict(torch.load(ckpt_path)) + model = model.cuda() + return model + + def wrap_model(): + model = GPT(**kwargs) + wrap_model = bmt.BMTrainModelWrapper(model) + bmt.load(wrap_model, ckpt_path) + return model + + def list_model(): + model = GPT(**kwargs) + list_model = bmt.BMTrainModelWrapper(model) + list_model.transformers = bmt.TransformerBlockList([m for m in list_model.transformers]) + bmt.load(list_model, ckpt_path) + return model + + def pipe_model(): + model = GPT(**kwargs) + pipe_model = bmt.BMTrainModelWrapper(model) + for m in pipe_model.transformers: + assert isinstance(m, bmt.CheckpointBlock) + pipe_model.transformers = bmt.PipelineTransformerBlockList([m for m in pipe_model.transformers]) + bmt.load(pipe_model, ckpt_path) + return model + + models = { + "torch": torch_model, + "wrapper": wrap_model, + "blocklist": list_model, + "pipelist": pipe_model, + } + loss_funcs = { + "bmt_entropy": bmt.loss.FusedCrossEntropy, + "torch_entropy": torch.nn.CrossEntropyLoss, + } + optimizers = { + "bmt_adam": bmt.optim.AdamOptimizer, + "bmt_adam_offload": bmt.optim.AdamOffloadOptimizer, + "torch_adam": torch.optim.Adam, + } + + ret = {} + def add_to_check_list(m, l, o): + key, value = train((m, models[m]), (l, loss_funcs[l]), (o, optimizers[o])) + ret[key] = value + + if test_fp16: + kwargs["dtype"] = torch.half + make_ref_ckpt() + add_to_check_list("torch", "bmt_entropy", "bmt_adam") + add_to_check_list("wrapper", "bmt_entropy", "bmt_adam") + add_to_check_list("blocklist", "bmt_entropy", "bmt_adam") + add_to_check_list("pipelist", "bmt_entropy", "bmt_adam") + add_to_check_list("blocklist", "torch_entropy", "bmt_adam") + add_to_check_list("blocklist", "bmt_entropy", "bmt_adam_offload") + if bmt.rank() == 0: + os.remove(ckpt_path) + check(ret) + + if test_fp32: + kwargs["dtype"] = torch.float + make_ref_ckpt() + add_to_check_list("torch", "torch_entropy", "bmt_adam") + add_to_check_list("wrapper", "torch_entropy", "bmt_adam") + add_to_check_list("blocklist", "torch_entropy", "bmt_adam") + add_to_check_list("pipelist", "torch_entropy", "bmt_adam") + add_to_check_list("blocklist", "torch_entropy", "bmt_adam_offload") + add_to_check_list("blocklist", "torch_entropy", "torch_adam") + if bmt.rank() == 0: + os.remove(ckpt_path) + check(ret) + +def check(ret): + if bmt.rank() == 0: + for k1, v1 in ret.items(): + for k2, v2 in ret.items(): + print(f"checking {k1} vs. {k2}") + check_param(v1[1], v2[1]) + bmt.synchronize() + ret.clear() + +def check_param(info1, info2): + for l1, l2 in zip(info1, info2): + for key in ['std', 'mean', 'max', 'min']: + v1 = l1[key] + v2 = l2[key] + assert_lt(abs(v1-v2), 1e-2) + +if __name__ == '__main__': + bmt.init_distributed(pipe_size=2) + + test_main(test_fp16=True, test_fp32=True) \ No newline at end of file diff --git a/tests/utils.py b/tests/utils.py new file mode 100644 index 00000000..62831fea --- /dev/null +++ b/tests/utils.py @@ -0,0 +1,14 @@ +def assert_eq(a, b): + assert a == b, f"{a} != {b}" + +def assert_neq(a, b): + assert a != b, f"{a} == {b}" + +def assert_lt(a, b): + assert a < b, f"{a} >= {b}" + +def assert_gt(a, b): + assert a > b, f"{a} <= {b}" + +def assert_all_eq(a, b): + assert_eq((a==b).all(), True) \ No newline at end of file