github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity
ColossalAI/tests/test_layers/test_2p5d/checks_2p5d/common.py

14 lines
220 B
Python
Raw Normal View History

Migrated project
2021-10-28 16:21:23 +00:00
import torch
TESSERACT_DIM = 2
TESSERACT_DEP = 2
BATCH_SIZE = 8
SEQ_LENGTH = 8
HIDDEN_SIZE = 8
moved env variables to global variables; (#215) added branch context; added vocab parallel layers; moved split_batch from load_batch to tensor parallel embedding layers; updated gpt model; updated unit test cases; fixed few collective communicator bugs
2022-02-14 03:15:02 +00:00
NUM_CLASSES = 8
VOCAB_SIZE = 16
IMG_SIZE = 16
Migrated project
2021-10-28 16:21:23 +00:00
def check_equal(A, B):
moved env variables to global variables; (#215) added branch context; added vocab parallel layers; moved split_batch from load_batch to tensor parallel embedding layers; updated gpt model; updated unit test cases; fixed few collective communicator bugs
2022-02-14 03:15:02 +00:00
assert torch.allclose(A, B, rtol=1e-5, atol=1e-2)