mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-09-20 09:01:06 +00:00
9ee197d0e9
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
17 lines
244 B
Python
17 lines
244 B
Python
#!/usr/bin/env python
|
|
# -*- encoding: utf-8 -*-
|
|
|
|
import torch
|
|
|
|
DEPTH = 2
|
|
BATCH_SIZE = 8
|
|
SEQ_LENGTH = 8
|
|
HIDDEN_SIZE = 8
|
|
NUM_CLASSES = 8
|
|
VOCAB_SIZE = 16
|
|
IMG_SIZE = 16
|
|
|
|
|
|
def check_equal(A, B):
|
|
assert torch.allclose(A, B, rtol=1e-3, atol=1e-2)
|