| |
|
|
| import functools |
|
|
| import torch |
| from packaging import version |
|
|
|
|
| def contiguous(fn): |
| @functools.wraps(fn) |
| def wrapper(ctx, *args, **kwargs): |
| return fn(ctx, |
| *(i if not isinstance(i, torch.Tensor) else i.contiguous() for i in args), |
| **{k: (v if not isinstance(v, torch.Tensor) else v.contiguous()) for k, v in kwargs.items()}) |
| return wrapper |
|
|
|
|
| def require_version(version, hint): |
| def decorator(fn): |
| @functools.wraps(fn) |
| def wrapper(ctx, *args, **kwargs): |
| from transformers.utils.versions import require_version |
| require_version(version, hint) |
| return fn(ctx, |
| *(i if not isinstance(i, torch.Tensor) else i.contiguous() for i in args), |
| **{k: (v if not isinstance(v, torch.Tensor) else v.contiguous()) for k, v in kwargs.items()}) |
| return wrapper |
| return decorator |
|
|
|
|
| def checkpoint(func): |
| def wrapper(*args, **kwargs): |
| return torch.utils.checkpoint.checkpoint(func, *args, **kwargs) |
| return wrapper |
|
|
|
|
| if version.parse(torch.__version__) >= version.parse("2.4"): |
| autocast_custom_fwd = functools.partial(torch.amp.custom_fwd, device_type="cuda") |
| autocast_custom_bwd = functools.partial(torch.amp.custom_bwd, device_type="cuda") |
| else: |
| autocast_custom_fwd = torch.cuda.amp.custom_fwd |
| autocast_custom_bwd = torch.cuda.amp.custom_bwd |
