forked from allenai/OLMo-core
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathMakefile
More file actions
130 lines (113 loc) · 4.48 KB
/
Copy pathMakefile
File metadata and controls
130 lines (113 loc) · 4.48 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
.PHONY : checks
checks : style-check lint-check type-check
.PHONY : style-check
style-check :
@echo "======== running isort... ========"
@isort --check .
@echo "======== running black... ========"
@black --check .
.PHONY : lint-check
lint-check :
@echo "======== running ruff... ========="
@ruff check .
.PHONY : type-check
type-check :
@echo "======== running mypy... ========="
@mypy src/
.PHONY : style
style:
@echo "======== formatting with isort... ========"
@isort .
@echo "======== formatting with black... ========"
@black .
.PHONY : docs
docs :
rm -rf docs/build/
sphinx-autobuild -b html --watch src/olmo_core/ --watch README.md docs/source/ docs/build/
.PHONY : build
build :
rm -rf *.egg-info/
python -m build
####################################################################################################
# Docker build
####################################################################################################
#-----------------#
# Build variables #
#-----------------#
# NOTE: When upgrading dependency versions (like for torch) make sure:
# * The corresponding versions specified in 'pyproject.toml' include the new version.
# * The versions installed in '.github/actions/setup-python-env/action.yml' match if necessary.
# NOTE: See https://hub.docker.com/r/nvidia/cuda/tags?name=devel-ubuntu22.04 for available CUDA versions.
CUDA_VERSION = 12.8.1
CUDA_VERSION_PATH=cu$(shell echo $(CUDA_VERSION) | cut -d"." -f1-2 | tr -d .)
PYTHON_VERSION = 3.12
TORCH_VERSION = 2.10.0
TORCH_VERSION_SHORT = $(shell echo $(TORCH_VERSION) | tr -d .)
INSTALL_CHANNEL = whl
GROUPED_GEMM_SHA = "f1429a3c44c98f7912aa4b00125144cdf4e7fdb2"
FLASH_ATTN_VERSION = 2.8.2
FLASH_ATTN_3_SHA = "060c9188beec3a8b62b33a3bfa6d5d2d44975fab"
# NOTE: At the moment FA4 isn't compatible with PyTorch 2.10, so to build FA4 you should set
# TORCH_VERSION to 2.9.1.
# FLASH_ATTN_4_SHA = "2580b5a4882562640f3cfbffd2bb8d2de9268f9f"
FLASH_ATTN_4_SHA = ""
FA3_MAX_JOBS = 64
TE_VERSION = 2.9
RING_FLASH_ATTN_VERSION = 0.1.8
LIGER_KERNEL_VERSION = 0.6.4
# NOTE: Quack currently requires CUDA 12.9 or higher and PyTorch 2.9.1
# QUACK_VERSION = 0.2.4
QUACK_VERSION = ""
#--------------#
# Build naming #
#--------------#
VERSION = $(shell python src/olmo_core/version.py)
VERSION_SHORT = $(shell python src/olmo_core/version.py short)
IMAGE_SUFFIX = $(shell date "+%Y-%m-%d")
IMAGE_TAG = tch$(TORCH_VERSION_SHORT)$(CUDA_VERSION_PATH)-$(IMAGE_SUFFIX)
.PHONY : docker-image
docker-image :
docker build -f src/Dockerfile \
--build-arg BUILDKIT_INLINE_CACHE=1 \
--build-arg CUDA_VERSION=$(CUDA_VERSION) \
--build-arg CUDA_VERSION_PATH=$(CUDA_VERSION_PATH) \
--build-arg PYTHON_VERSION=$(PYTHON_VERSION) \
--build-arg TORCH_VERSION=$(TORCH_VERSION) \
--build-arg INSTALL_CHANNEL=$(INSTALL_CHANNEL) \
--build-arg GROUPED_GEMM_SHA=$(GROUPED_GEMM_SHA) \
--build-arg FLASH_ATTN_VERSION=$(FLASH_ATTN_VERSION) \
--build-arg FLASH_ATTN_3_SHA=$(FLASH_ATTN_3_SHA) \
--build-arg FLASH_ATTN_4_SHA=$(FLASH_ATTN_4_SHA) \
--build-arg FA3_MAX_JOBS=$(FA3_MAX_JOBS) \
--build-arg TE_VERSION=$(TE_VERSION) \
--build-arg RING_FLASH_ATTN_VERSION=$(RING_FLASH_ATTN_VERSION) \
--build-arg LIGER_KERNEL_VERSION=$(LIGER_KERNEL_VERSION) \
--build-arg QUACK_VERSION=$(QUACK_VERSION) \
--target release \
-t olmo-core:$(IMAGE_TAG) .
@docker run --rm olmo-core:$(IMAGE_TAG) python -c \
'import torch; import transformer_engine.pytorch; import flash_attn; import flash_attn_3.flash_attn_interface'
@echo "✓ Image validated. Python environment:"
@echo ""
@docker run --rm olmo-core:$(IMAGE_TAG) pip list
@echo ""
@echo "✓ Build complete: olmo-core:$(IMAGE_TAG) (size=$$(docker inspect -f '{{ .Size }}' olmo-core:$(IMAGE_TAG) | numfmt --to=si))"
@echo ""
.PHONY : ghcr-image
ghcr-image : docker-image
docker tag olmo-core:$(IMAGE_TAG) ghcr.io/allenai/olmo-core:$(IMAGE_TAG)
docker push ghcr.io/allenai/olmo-core:$(IMAGE_TAG)
docker tag olmo-core:$(IMAGE_TAG) ghcr.io/allenai/olmo-core:latest
docker push ghcr.io/allenai/olmo-core:latest
BEAKER_WORKSPACE = ai2/OLMo-core
BEAKER_USER = $(shell beaker account whoami --format=json | jq -r '.[0].name')
.PHONY : beaker-image
beaker-image : docker-image
@./src/scripts/beaker/create_beaker_image.sh olmo-core:$(IMAGE_TAG) olmo-core-$(IMAGE_TAG) $(BEAKER_WORKSPACE)
@echo "✓ Done"
.PHONY : get-beaker-workspace
get-beaker-workspace :
@echo $(BEAKER_WORKSPACE)
.PHONY : get-full-beaker-image-name
get-full-beaker-image-name :
@./src/scripts/beaker/get_full_image_name.sh $(IMAGE_TAG) $(BEAKER_WORKSPACE)