Buckets:
| """Check the allocated lab GPUs and the FlashAttention forward/backward kernel.""" | |
| import json | |
| import os | |
| import torch | |
| import flash_attn | |
| from flash_attn import flash_attn_func | |
| def main(): | |
| expected = len(os.environ["BAGEL_ALLOCATED_GPU_IDS"].split(",")) | |
| if torch.cuda.device_count() != expected: | |
| raise RuntimeError("GPU queue allocation differs from visible device count") | |
| if not torch.__version__.startswith("2.7.1"): | |
| raise RuntimeError(f"unexpected Torch version: {torch.__version__}") | |
| records = [] | |
| for index in range(expected): | |
| with torch.cuda.device(index): | |
| q = torch.randn(1, 16, 4, 64, device=f"cuda:{index}", | |
| dtype=torch.bfloat16, requires_grad=True) | |
| result = flash_attn_func(q, q, q, causal=True) | |
| result.float().square().mean().backward() | |
| assert torch.isfinite(result).all() and torch.isfinite(q.grad).all() | |
| torch.cuda.synchronize() | |
| records.append({"gpu": index, "name": torch.cuda.get_device_name(index), | |
| "capability": torch.cuda.get_device_capability(index)}) | |
| print(json.dumps({"torch": torch.__version__, "cuda": torch.version.cuda, | |
| "flash_attention": flash_attn.__version__, "devices": records})) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 1.35 kB
- Xet hash:
- 94c8210641e4848d777bdf09e86e6b6752dfdf0049143a268484421f3fa49869
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.