python-megatron-core
CLEAN
maintainer Smoolak
0 votes
scanned 2026-09-21 00:26:32.109917
Triggered rules
Clean
AI review downgraded a static finding
llm_review
The static rules flagged this LOW, but an AI model (qwen/qwen3-235b-a22b-2507) reviewed the full PKGBUILD and judged it CLEAN (confidence 95%): The package builds from a verifiable source on NVIDIA's official GitHub repository, uses standard Python build tools, and includes reasonable checks; the low-vote and recent-upload flags are common for legitimate new AUR packages.
1 higher static finding superseded - not the current verdict (shown for transparency)
Low
Few votes, recently uploaded
zero_votes_recent
Uploaded within the last 14 days with 2 or fewer community votes — little peer review so far.
PKGBUILD
1
# Maintainer: Smoolak <smoolak@gmail.com>
2
3
pkgname=python-megatron-core
4
_pkgname=Megatron-LM
5
pkgver=0.19.2
6
pkgrel=1
7
_archive="$_pkgname-core_v$pkgver"
8
pkgdesc='NVIDIA Megatron-Core: GPU-optimized building blocks for training transformer models at scale'
9
arch=('x86_64')
10
url='https://github.com/NVIDIA/Megatron-LM'
11
license=('Apache-2.0' 'BSD-3-Clause' 'MIT')
12
depends=(
13
'python'
14
'python-numpy'
15
'python-packaging>=24.2'
16
'python-pytorch>=2.6.0'
17
'python-setuptools>=80'
18
'python-triton'
19
'python-typing_extensions'
20
'python-yaml'
21
)
22
makedepends=(
23
'pybind11'
24
'python-build'
25
'python-installer'
26
'python-wheel'
27
)
28
checkdepends=(
29
'python-click'
30
'python-msgpack'
31
'python-omegaconf'
32
'python-pytest'
33
'python-requests'
34
'python-transformers'
35
)
36
optdepends=(
37
'python-absl: named communication-grid logging'
38
'python-accelerate: Hugging Face distributed training integration'
39
'python-apex: fused optimizers, normalization, and transformer kernels'
40
'python-boto3: S3-backed datasets'
41
'python-botocore: low-level S3 object-storage integration'
42
'python-causal-conv1d: state-space model convolution kernels'
43
'python-datasets: supervised fine-tuning datasets'
44
'python-einops: tensor rearrangements and multi-latent attention'
45
'python-flash-attn: FlashAttention support'
46
'python-flashinfer: accelerated inference kernels'
47
'python-flask: legacy text-generation server'
48
'python-flask-restful: legacy text-generation REST API'
49
'hypercorn: dynamic text-generation ASGI server'
50
'python-mamba-ssm: Mamba state-space models'
51
'python-msgpack: dynamic-inference request serialization'
52
'python-nvidia-ml-py: GPU energy monitoring'
53
'python-nvidia-modelopt: quantization and post-training optimization'
54
'python-nvshmem: NVSHMEM resharding and copy services'
55
'python-omegaconf: structured training configuration'
56
'python-orjson: fast OpenAI-compatible response serialization'
57
'python-psutil: asynchronous checkpoint process monitoring'
58
'python-pyzmq: dynamic-inference communication'
59
'python-quart: dynamic text-generation HTTP API'
60
'python-sentencepiece: SentencePiece tokenizers'
61
'python-tensorrt-llm: TensorRT-LLM engine export'
62
'python-tiktoken: TikToken tokenizers'
63
'python-tqdm: progress reporting'
64
'python-transformer-engine: FP8 training and fused transformer kernels'
65
'python-transformers: Hugging Face model and tokenizer interoperability'
66
'python-wandb: experiment logging'
67
'python-wget: Hugging Face tokenizer downloads'
68
)
69
source=("$pkgname-$pkgver.tar.gz::$url/archive/refs/tags/core_v$pkgver.tar.gz")
70
sha256sums=('4c16ccc2a0d4c1b6c96f59baaab99316ac3b39eca2ef960ddf4436a2962eb69a')
71
72
build() {
73
cd "$_archive"
74
export NO_VCS_VERSION=1
75
python -m build --wheel --no-isolation
76
}
77
78
check() {
79
cd "$_archive"
80
local _site_packages
81
_site_packages="$(python -c 'import sysconfig; print(sysconfig.get_path("purelib"))')"
82
83
rm -rf "$srcdir/test-root" "$srcdir/test-run" "$srcdir/test-data"
84
python -m installer --destdir="$srcdir/test-root" dist/*.whl
85
install -d "$srcdir/test-run" "$srcdir/test-data"
86
touch "$srcdir/test-data/.offline"
87
cp -a tests "$srcdir/test-run/"
88
sed -i \
89
's#data_path = Path("/opt/data")#data_path = Path(os.environ["MEGATRON_TEST_DATA"])#' \
90
"$srcdir/test-run/tests/unit_tests/conftest.py"
91
cd "$srcdir/test-run"
92
93
# Upstream's complete unit runner launches eight CUDA workers with NCCL.
94
# Run its broad CPU-only surface here. The remaining deselection needs
95
# optional Multi-Storage Client (not available in Arch/AUR); Transformers
96
# cases run with their declared check dependency. CUDA/NCCL/FP8 and
97
# multi-GPU suites remain hardware-bound.
98
CUDA_VISIBLE_DEVICES= \
99
HIP_VISIBLE_DEVICES= \
100
ROCR_VISIBLE_DEVICES= \
101
MEGATRON_TEST_DATA="$srcdir/test-data" \
102
NO_VCS_VERSION=1 \
103
PYTHONPATH="$srcdir/test-root$_site_packages" \
104
python -m pytest -q -o addopts='' --basetemp="$srcdir/pytest-tmp" \
105
tests/unit_tests/inference/test_common_inference_params.py \
106
tests/unit_tests/inference/test_inference_config.py \
107
tests/unit_tests/inference/test_inference_request.py \
108
tests/unit_tests/inference/test_scheduler.py \
109
tests/unit_tests/inference/test_stop_words.py \
110
tests/unit_tests/optimizer/test_optimizer_config.py \
111
tests/unit_tests/pipeline_parallel/test_helpers.py \
112
tests/unit_tests/resharding/test_communication_scheduler.py \
113
tests/unit_tests/resharding/test_dp_balancing.py \
114
tests/unit_tests/resharding/test_planner.py \
115
tests/unit_tests/resharding/test_task_segmenter.py \
116
tests/unit_tests/resharding/test_workload_packer.py \
117
tests/unit_tests/ssm/test_hybrid_layer_allocation.py \
118
tests/unit_tests/test_argument_utils.py \
119
tests/unit_tests/test_basic.py \
120
tests/unit_tests/test_model_configs.py \
121
tests/unit_tests/test_num_microbatches_calculator.py \
122
tests/unit_tests/test_num_floating_point_operations.py \
123
tests/unit_tests/test_optimizer_param_scheduler.py \
124
tests/unit_tests/test_start_eval_at_iter.py \
125
tests/unit_tests/test_typed_torch.py \
126
tests/unit_tests/training/config/test_container_base.py \
127
tests/unit_tests/training/config/test_instantiate_utils.py \
128
tests/unit_tests/training/config/test_target_allowlist.py \
129
tests/unit_tests/training/config/test_utils.py \
130
tests/unit_tests/training/config/test_yaml_utils.py \
131
tests/unit_tests/training/models/test_base.py \
132
tests/unit_tests/transformer/test_mup.py \
133
tests/unit_tests/transformer/test_quantization_config.py \
134
tests/unit_tests/transformer/test_spec_utils.py \
135
tests/unit_tests/utils/test_experimental_log_once.py \
136
tests/unit_tests/utils/test_log_utils.py \
137
tests/unit_tests/utils/test_slurm_utils.py \
138
--deselect=tests/unit_tests/training/config/test_container_base.py::TestConfigContainer_ToYaml::test_to_yaml_with_msc_url
139
140
# Exercise the compiled dataset helper and a single-rank, CPU-initialized GPT
141
# training slice: embeddings, vocabulary projection, loss, and gradients.
142
CUDA_VISIBLE_DEVICES= \
143
HIP_VISIBLE_DEVICES= \
144
ROCR_VISIBLE_DEVICES= \
145
NO_VCS_VERSION=1 \
146
PYTHONPATH="$srcdir/test-root$_site_packages" \
147
python - <<'PY'
148
import tempfile
149
150
import numpy as np
151
import torch
152
import torch.distributed as dist
153
import torch.nn.functional as F
154
155
from megatron.core import parallel_state
156
from megatron.core.datasets.helpers import build_sample_idx
157
from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_local_spec
158
from megatron.core.models.gpt.gpt_model import GPTModel
159
from megatron.core.transformer.transformer_config import TransformerConfig
160
161
sample_idx = build_sample_idx(
162
np.array([5, 7], dtype=np.int32),
163
np.array([0, 1], dtype=np.int32),
164
sequence_length=4,
165
num_epochs=1,
166
tokens_per_epoch=12,
167
)
168
np.testing.assert_array_equal(sample_idx, [[0, 0], [0, 4], [1, 3]])
169
170
with tempfile.TemporaryDirectory() as tmpdir:
171
dist.init_process_group(
172
"gloo", rank=0, world_size=1, init_method=f"file://{tmpdir}/store"
173
)
174
parallel_state.initialize_model_parallel(
175
tensor_model_parallel_size=1, pipeline_model_parallel_size=1
176
)
177
try:
178
config = TransformerConfig(
179
num_layers=2,
180
hidden_size=32,
181
num_attention_heads=4,
182
ffn_hidden_size=64,
183
use_cpu_initialization=True,
184
hidden_dropout=0.0,
185
attention_dropout=0.0,
186
)
187
torch.manual_seed(1234)
188
model = GPTModel(
189
config=config,
190
transformer_layer_spec=get_gpt_layer_local_spec(),
191
vocab_size=64,
192
max_sequence_length=8,
193
position_embedding_type="learned_absolute",
194
parallel_output=False,
195
)
196
input_ids = torch.tensor([[1, 2, 3, 4], [4, 3, 2, 1]])
197
position_ids = torch.arange(4).unsqueeze(0).expand(2, -1)
198
hidden = model.embedding(input_ids, position_ids)
199
logits, bias = model.output_layer(hidden, weight=None)
200
targets = torch.tensor([[2, 3, 4, 5], [3, 2, 1, 0]]).T
201
loss = F.cross_entropy(logits.reshape(-1, 64), targets.reshape(-1))
202
loss.backward()
203
204
assert bias is None
205
assert logits.shape == (4, 2, 64)
206
assert torch.isfinite(logits).all()
207
assert 3.0 < loss.item() < 5.5
208
assert model.embedding.word_embeddings.weight.grad.norm() > 0
209
assert model.output_layer.weight.grad.norm() > 0
210
finally:
211
parallel_state.destroy_model_parallel()
212
dist.destroy_process_group()
213
PY
214
}
215
216
package() {
217
cd "$_archive"
218
python -m installer --destdir="$pkgdir" dist/*.whl
219
install -Dm644 LICENSE "$pkgdir/usr/share/licenses/$pkgname/LICENSE"
220
sed -n '/^This repository also contains code/,$p' LICENSE \
221
> "$pkgdir/usr/share/licenses/$pkgname/THIRD-PARTY-LICENSES"
222
}
223
Scan history
| Scanned at (UTC) | Severity | Rules |
|---|---|---|
| 2026-09-21 00:26:32 | Clean | 2 |
| 2026-09-20 23:33:19 | Low | 1 |