|
|
|
|
|
|
|
|
|
|
|
import bisect |
|
|
import inspect |
|
|
import os |
|
|
import random |
|
|
from contextlib import contextmanager |
|
|
from datetime import datetime |
|
|
from typing import Union |
|
|
|
|
|
import numpy as np |
|
|
import torch |
|
|
|
|
|
import internlm |
|
|
|
|
|
CURRENT_TIME = None |
|
|
|
|
|
|
|
|
def parse_args(): |
|
|
parser = internlm.get_default_parser() |
|
|
args = parser.parse_args() |
|
|
|
|
|
return args |
|
|
|
|
|
|
|
|
def get_master_node(): |
|
|
import subprocess |
|
|
|
|
|
if os.getenv("SLURM_JOB_ID") is None: |
|
|
raise RuntimeError("get_master_node can only used in Slurm launch!") |
|
|
result = subprocess.check_output('scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1', shell=True) |
|
|
result = result.decode("utf8").strip() |
|
|
return result |
|
|
|
|
|
|
|
|
def move_norm_to_cuda(norm: Union[float, torch.Tensor]) -> Union[float, torch.Tensor]: |
|
|
if torch.is_tensor(norm) and norm.device.type != "cuda": |
|
|
norm = norm.to(torch.cuda.current_device()) |
|
|
return norm |
|
|
|
|
|
|
|
|
def _move_tensor(element): |
|
|
if not torch.is_tensor(element): |
|
|
|
|
|
for item in element: |
|
|
if isinstance(item, dict): |
|
|
for key, value in item.items(): |
|
|
assert not value.is_cuda, "elements are already on devices." |
|
|
item[key] = value.to(get_current_device()).detach() |
|
|
elif isinstance(item, list): |
|
|
for index, value in enumerate(item): |
|
|
assert not value.is_cuda, "elements are already on devices." |
|
|
item[index] = value.to(get_current_device()).detach() |
|
|
elif torch.is_tensor(item): |
|
|
if not item.is_cuda: |
|
|
item = item.to(get_current_device()).detach() |
|
|
else: |
|
|
assert torch.is_tensor(element), f"element should be of type tensor, but got {type(element)}" |
|
|
if not element.is_cuda: |
|
|
element = element.to(get_current_device()).detach() |
|
|
return element |
|
|
|
|
|
|
|
|
def move_to_device(data): |
|
|
if isinstance(data, torch.Tensor): |
|
|
data = data.to(get_current_device()) |
|
|
elif isinstance(data, (list, tuple)): |
|
|
data_to_return = [] |
|
|
for element in data: |
|
|
if isinstance(element, dict): |
|
|
data_to_return.append({k: _move_tensor(v) for k, v in element.items()}) |
|
|
else: |
|
|
data_to_return.append(_move_tensor(element)) |
|
|
data = data_to_return |
|
|
elif isinstance(data, dict): |
|
|
data = {k: _move_tensor(v) for k, v in data.items()} |
|
|
else: |
|
|
raise TypeError(f"Expected batch data to be of type torch.Tensor, list, tuple, or dict, but got {type(data)}") |
|
|
return data |
|
|
|
|
|
|
|
|
def get_tensor_norm(norm: Union[float, torch.Tensor], move_to_cuda) -> torch.Tensor: |
|
|
if isinstance(norm, float): |
|
|
norm = torch.Tensor([norm]) |
|
|
if move_to_cuda: |
|
|
norm = norm.to(torch.cuda.current_device()) |
|
|
return norm |
|
|
|
|
|
|
|
|
def get_current_device() -> torch.device: |
|
|
""" |
|
|
Returns currently selected device (gpu/cpu). |
|
|
If cuda available, return gpu, otherwise return cpu. |
|
|
""" |
|
|
if torch.cuda.is_available(): |
|
|
return torch.device(f"cuda:{torch.cuda.current_device()}") |
|
|
else: |
|
|
return torch.device("cpu") |
|
|
|
|
|
|
|
|
def get_batch_size(data): |
|
|
if isinstance(data, torch.Tensor): |
|
|
return data.size(0) |
|
|
elif isinstance(data, (list, tuple)): |
|
|
if isinstance(data[0], dict): |
|
|
return data[0][list(data[0].keys())[0]].size(0) |
|
|
return data[0].size(0) |
|
|
elif isinstance(data, dict): |
|
|
return data[list(data.keys())[0]].size(0) |
|
|
|
|
|
|
|
|
def filter_kwargs(func, kwargs): |
|
|
sig = inspect.signature(func) |
|
|
return {k: v for k, v in kwargs.items() if k in sig.parameters} |
|
|
|
|
|
|
|
|
def launch_time(): |
|
|
global CURRENT_TIME |
|
|
if not CURRENT_TIME: |
|
|
CURRENT_TIME = datetime.now().strftime("%b%d_%H-%M-%S") |
|
|
return CURRENT_TIME |
|
|
|
|
|
|
|
|
def set_random_seed(seed): |
|
|
"""Set random seed for reproducability.""" |
|
|
|
|
|
if seed is not None: |
|
|
assert seed > 0 |
|
|
random.seed(seed) |
|
|
np.random.seed(seed) |
|
|
torch.manual_seed(seed) |
|
|
torch.cuda.manual_seed(seed) |
|
|
|
|
|
torch.cuda.manual_seed_all(seed) |
|
|
|
|
|
|
|
|
@contextmanager |
|
|
def conditional_context(context_manager, enable=True): |
|
|
if enable: |
|
|
with context_manager: |
|
|
yield |
|
|
else: |
|
|
yield |
|
|
|
|
|
|
|
|
class BatchSkipper: |
|
|
""" |
|
|
BatchSkipper is used to determine whether to skip the current batch_idx. |
|
|
""" |
|
|
|
|
|
def __init__(self, skip_batches): |
|
|
if skip_batches == "": |
|
|
pass |
|
|
intervals = skip_batches.split(",") |
|
|
spans = [] |
|
|
if skip_batches != "": |
|
|
for interval in intervals: |
|
|
if "-" in interval: |
|
|
start, end = map(int, interval.split("-")) |
|
|
else: |
|
|
start, end = int(interval), int(interval) |
|
|
if spans: |
|
|
assert spans[-1] <= start |
|
|
spans.extend((start, end + 1)) |
|
|
self.spans = spans |
|
|
|
|
|
def __call__(self, batch_count): |
|
|
index = bisect.bisect_right(self.spans, batch_count) |
|
|
return index % 2 == 1 |
|
|
|
|
|
|
|
|
class SingletonMeta(type): |
|
|
""" |
|
|
Singleton Meta. |
|
|
""" |
|
|
|
|
|
_instances = {} |
|
|
|
|
|
def __call__(cls, *args, **kwargs): |
|
|
if cls not in cls._instances: |
|
|
cls._instances[cls] = super().__call__(*args, **kwargs) |
|
|
else: |
|
|
assert ( |
|
|
len(args) == 0 and len(kwargs) == 0 |
|
|
), f"{cls.__name__} is a singleton class and a instance has been created." |
|
|
return cls._instances[cls] |
|
|
|
|
|
|
|
|
def get_megatron_flops( |
|
|
elapsed_time_per_iter, |
|
|
checkpoint=False, |
|
|
seq_len=2048, |
|
|
hidden_size=12, |
|
|
num_layers=32, |
|
|
vocab_size=12, |
|
|
global_batch_size=4, |
|
|
global_world_size=1, |
|
|
mlp_ratio=4, |
|
|
use_swiglu=True, |
|
|
): |
|
|
""" |
|
|
Calc flops based on the paper of Megatron https://deepakn94.github.io/assets/papers/megatron-sc21.pdf |
|
|
""" |
|
|
|
|
|
checkpoint_activations_factor = 4 if checkpoint else 3 |
|
|
|
|
|
if use_swiglu: |
|
|
mlp_ratio = mlp_ratio * 3 / 2 |
|
|
|
|
|
flops_per_iteration = ( |
|
|
checkpoint_activations_factor |
|
|
* ( |
|
|
(8 + mlp_ratio * 4) * global_batch_size * seq_len * hidden_size**2 |
|
|
+ 4 * global_batch_size * seq_len**2 * hidden_size |
|
|
) |
|
|
) * num_layers + 6 * global_batch_size * seq_len * hidden_size * vocab_size |
|
|
|
|
|
tflops = flops_per_iteration / (elapsed_time_per_iter * global_world_size * (10**12)) |
|
|
return tflops |
|
|
|
|
|
|
|
|
class DummyProfile: |
|
|
""" |
|
|
Dummy Profile. |
|
|
""" |
|
|
|
|
|
def __init__(self, *args, **kwargs) -> None: |
|
|
pass |
|
|
|
|
|
def __enter__(self): |
|
|
return self |
|
|
|
|
|
def __exit__(self, a, b, c): |
|
|
pass |
|
|
|
|
|
def step(self): |
|
|
pass |
|
|
|
