diff --git a/.alignment.sh.swp b/.alignment.sh.swp
new file mode 100644
index 00000000..e4fe36cd
Binary files /dev/null and b/.alignment.sh.swp differ
diff --git a/.gitattributes b/.gitattributes
new file mode 100644
index 00000000..f0ee6826
--- /dev/null
+++ b/.gitattributes
@@ -0,0 +1,4 @@
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+checkpoints/LLaVA-OneVision-1.5-4B-stage0/** filter=lfs diff=lfs merge=lfs -text
diff --git a/.gitignore b/.gitignore
index 00186b4c..2220df92 100644
--- a/.gitignore
+++ b/.gitignore
@@ -60,3 +60,20 @@ coverage.xml
**/.ipynb_checkpoints/
.vscode/
checkpoints/
+stage_1_alignment_llava_ov_4b/
+tmp/
+Stage1/training_output.log
+
+# Local training/evaluation artifacts
+*.log
+*.out
+wandb/
+tensorboard/
+eval_outputs/
+data/hf_cache/
+data/hf_cache_midtraining_stream/
+data/midtraining_*_webdataset/
+stage_1_5_midtraining_*/
+stage_1_alignment_mobilellm_140m/
+stage_1_alignment_mobilellm_140m_fastvlm_faithful_smoke/
+stage_1_alignment_mobilellm_140m_fastvlm_faithful/
diff --git a/dockerfile b/Dockerfile
similarity index 100%
rename from dockerfile
rename to Dockerfile
diff --git a/README.md b/README.md
index 2952e9d0..abdd0025 100644
--- a/README.md
+++ b/README.md
@@ -1,494 +1,577 @@
-
-
-
-
-
-
+# LLaVA-OneVision-1.5 ร Mobile-LLM Integration
+
+> **Branch:** `mobile-llm-integration`
+>
+> This branch adapts the [LLaVA-OneVision-1.5](https://github.com/EvolvingLMMs-Lab/LLaVA-OneVision-1.5) training framework to a fully **mobile-optimized** multimodal pipeline by replacing the original vision encoder and language model with lightweight alternatives:
+>
+> | Component | Original (upstream) | This branch |
+> |---|---|---|
+> | Vision Encoder | RICE ViT-Large (560 px) | **FastViT / MobileCLIP-L** (1024 px, Apple ml-fastvlm) |
+> | Language Model | Qwen3-4B | **MobileLLM-R1-140M** (Facebook, 140M params) |
+> | Adapter | 2-layer MLP | 2-layer MLP (3072 โ 576, re-initialized) |
+> | Training Stage shown | Stage 1 alignment | **Stage 1 alignment** (adapter only frozen: vision + LLM) |
-
- Fully Open Framework for Democratized Multimodal Training
-
+---
+## Table of Contents
+- [Architecture Overview](#architecture-overview)
+- [Repository Structure](#repository-structure)
+- [Dependencies & Installation](#dependencies--installation)
+- [Downloading Pretrained Checkpoints](#downloading-pretrained-checkpoints)
+- [Downloading the Dataset](#downloading-the-dataset)
+- [Running Stage 1 Alignment Training](#running-stage-1-alignment-training)
+- [Demo: End-to-End Inference](#demo-end-to-end-inference)
+- [Training Logs & Results](#training-logs--results)
+- [Modifications vs. Upstream](#modifications-vs-upstream)
+- [Credits & Citations](#credits--citations)
-
+---
-๐ค **[Models and Datasets](https://huggingface.co/collections/lmms-lab/llava-onevision-15-68d385fe73b50bd22de23713)** |
-๐ฅ๏ธ **[Demo](https://huggingface.co/spaces/lmms-lab/LLaVA-OneVision-1.5)** |
-๐ **[Technical Report](https://arxiv.org/abs/2509.23661)** |
-๐ฐ **[Zhihu](https://www.zhihu.com/question/1959577143697707446)** |
-๐ **[Xiaohongshu](http://xhslink.com/o/4nXL6EXDTqv)**
+## Architecture Overview
-
+```
+โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+โ Mobile Multimodal Pipeline โ
+โ โ
+โ Image (1024ร1024) โ
+โ โ โ
+โ โผ โ
+โ โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ โ
+โ โ FastViT MobileCLIP-L [FROZEN] โ โ
+โ โ โข Apple ml-fastvlm architecture โ โ
+โ โ โข 1024ร1024 input, 64ร64 patch grid โ โ
+โ โ โข Output: [B, 256 tokens, 3072 dim] โ โ
+โ โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ โ
+โ โ โ
+โ โผ (32 images packed โ flattened to [8192, 3072]) โ
+โ โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ โ
+โ โ 2-Layer MLP Adapter [TRAINABLE] โ โ
+โ โ โข Linear(3072, 3072) + GELU โ โ
+โ โ โข Linear(3072, 576) โ โ
+โ โ โข Output: [8192, 576] โ โ
+โ โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ โ
+โ โ โ
+โ โผ (merged into text token sequence at <|image_pad|> slots) โ
+โ โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ โ
+โ โ MobileLLM-R1-140M [FROZEN] โ โ
+โ โ โข 15 transformer layers โ โ
+โ โ โข Hidden: 576, Heads: 9, KV-heads: 3 โ โ
+โ โ โข FFN: 2048 (SwiGLU), 32k context โ โ
+โ โ โข QK LayerNorm enabled โ โ
+โ โ โข GQA (9 query / 3 key-value heads) โ โ
+โ โ โข Output: per-token cross-entropy loss โ โ
+โ โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ โ
+โ โ
+โ Stage 1: Only the Adapter is trained (โ 3.5M params) โ
+โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+```
+
+### Pipeline Step-by-Step (one forward pass)
+
+```
+[1/6] BATCH
+ input_ids : (1, 1909) torch.int64
+ images : (32, 3, 1024, 1024) torch.bfloat16
+ labels : (1, 1909) torch.int64
+ loss_mask : 532 / 1909 tokens (27.9%) contribute to loss
+
+[2/6] VISION ENCODER [FastViT MobileCLIP-L | FROZEN]
+ in : (32, 3, 1024, 1024) torch.bfloat16
+ out : (32, 256, 3072) torch.bfloat16 grad=False
+
+[3/6] ADAPTER [2-layer MLP 3072โ576 | TRAINABLE]
+ in : (32, 256, 3072)
+ out : (32, 256, 576) grad=True
+
+[4/6] TOKEN FUSION
+ text embeddings : (1909, 1, 576) torch.bfloat16 grad=False
+ image slots : 32 tokens replaced with vision embeddings
+ combined : (1909, 1, 576) grad=True
+
+[5/6] LANGUAGE MODEL [MobileLLM-R1-140M | FROZEN]
+ in : (1909, 1, 576) torch.bfloat16
+ out : (1, 1909) torch.float32 grad=True
+
+[6/6] LOSS / LOGITS
+ loss (mean over 532 tokens) : 11.7835 โ step 1 (random adapter init)
+ top-1 accuracy : 1/532 = 0.2% โ expected near-random at step 1
+```
---
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
+## Repository Structure
+
+```
+LLaVA-OneVision-1.5/
+โ
+โโโ README.md โ This file
+โโโ inference_fastvlm.py โ Demo: end-to-end inference script
+โ
+โโโ examples/llava_ov_1_5/
+โ โโโ quick_start/
+โ โโโ stage_1_alignment_mobilellm_140m.sh โ Main training launcher
+โ
+โโโ stage_1_alignment_mobilellm_140m/ โ Training outputs (this run)
+โ โโโ iter_0000500/ โ Latest checkpoint (500 steps)
+โ โ โโโ mp_rank_00/
+โ โ โโโ model_optim_rng.pt โ Model + optimizer state (529 MB, Git LFS)
+โ โ โโโ distrib_optim.pt โ Distributed optimizer state (129 MB, Git LFS)
+โ โโโ latest_checkpointed_iteration.txt
+โ โโโ run_2026-04-29_13:33:58_*.log โ Full training log (500 steps)
+โ
+โโโ aiak_training_llm/
+โ โโโ models/
+โ โ โโโ llavaov_1_5/ โ Core model: integrates all three components
+โ โ โ โโโ llavaov_1_5_model.py โ Main forward pass with 6-step pipeline logging
+โ โ โ โโโ llavaov_1_5_config.py โ Model configuration (sets qk_layernorm=True)
+โ โ โ โโโ llavaov_1_5_layer_spec.pyโ Transformer layer spec (imports from mobilellm)
+โ โ โ โโโ llavaov_1_5_provider.py โ Model provider for Megatron training loop
+โ โ โ โโโ rice_vision_model.py โ Vision model base (rotary pos emb for vision)
+โ โ โ
+โ โ โโโ mobilellm/ โ MobileLLM-R1-140M integration
+โ โ โ โโโ mobilellm_model.py โ Megatron GPT model wrapper for MobileLLM
+โ โ โ โโโ mobilellm_config.py โ Reads HF config.json โ TransformerConfig
+โ โ โ โโโ mobilellm_layer_spec.py โ Layer spec with correct QK-norm handling
+โ โ โ โโโ mobilellm_provider.py โ Model provider
+โ โ โ
+โ โ โโโ fastvit/ โ FastViT / MobileCLIP-L vision encoder
+โ โ โโโ fastvit_vision_model.py โ Megatron-compatible wrapper
+โ โ โโโ mobileclip_encoder.py โ MobileCLIPVisionTower (loads pretrained weights)
+โ โ โโโ fastvit_preprocessor.py โ Image preprocessing (pad to square, resize)
+โ โ โโโ mm_utils.py โ expand2square, image padding utilities
+โ โ โโโ mobileclip/ โ Apple ml-fastvlm model code (FastViT backbone)
+โ โ โโโ mci.py โ FastViT model class
+โ โ โโโ ...
+โ โ
+โ โโโ data/multimodal/
+โ โ โโโ qwen2vl_task_encoder.py โ Main data pipeline (FastViT path added)
+โ โ โโโ task_encoder.py โ Base task encoder
+โ โ
+โ โโโ train/
+โ โโโ pretrain/
+โ โ โโโ pretrain_llavaov_1_5.py โ forward_step, loss_func, training entry
+โ โโโ training_utils.py โ Megatron training loop
+โ
+โโโ aiak_megatron/megatron/core/
+โ โโโ transformer/
+โ โ โโโ attention.py โ Patched: debug prints removed
+โ โ โโโ dot_product_attention.py โ Patched: debug prints removed, GQA shape fix
+โ โโโ extensions/
+โ โโโ transformer_engine.py โ Patched: ROCm/AMD compatibility
+โ โโโ transformer_engine2.py โ Patched: ROCm/AMD compatibility
+โ
+โโโ apex/csrc/
+โ โโโ mlp.cpp โ Patched: PyTorch 2.x API (.scalar_type())
+โ โโโ fused_dense.cpp โ Patched: PyTorch 2.x API (.scalar_type())
+โ
+โโโ checkpoints/
+ โโโ mobilellm-fastvit-merged-tp1-pp1/โ Pretrained merged checkpoint (NOT in git)
+```
---
+## Dependencies & Installation
-## NEWS
-- 2025-09-30: Released the [Offline Data Packing Guide](examples_offline_packing).
-- 2025-09-30: Released the LLaVA-OneVision-1.5 [Technical Report](https://arxiv.org/abs/2509.23661).
+### Hardware Requirements
+- **GPU:** NVIDIA GPU with โฅ 16 GB VRAM (tested on A100 80 GB)
+ *Note:* This codebase was also patched for ROCm/AMD compatibility.
+- **CUDA:** 12.x (tested with CUDA 12.1)
+- **RAM:** โฅ 32 GB system RAM recommended
-## Contents
-
-- [Introduction](#introduction)
-- [Models](#models)
-- [Datasets](#datasets)
-- [Results](#evaluation-results)
-- [Quick Start with Hugging Face](#quick-start-with-huggingface)
-- [Evaluation](#evaluation)
-- [Quick Start For Training](#quick-start-guide)
-- [Fully Reproducing Guide](#fully-reproducing-guide)
-- [Citation](#citation)
-- [Acknowledgement](#acknowledgement)
+### 1. Clone the repository
+```bash
+git clone https://github.com/RanaZay/LLaVA-OneVision-1.5.git
+cd LLaVA-OneVision-1.5
+git checkout mobile-llm-integration
+```
-## Introduction
-**LLaVA-OneVision-1.5** introduces a family of fully open-source large multimodal models (LMMs) that operate on **native-resolution images**, achieve **state-of-the-art** performance, and require comparatively **lower training costs**.
+### 2. Create conda environment
-#### **Superior Performance**
- - The model leads on multiple multimodal benchmarks and generally surpasses Qwen2.5-VL.
- - Training on native-resolution images significantly improves its visual understanding.
+```bash
+conda create -n llava-mobile python=3.10 -y
+conda activate llava-mobile
+```
-#### **High-Quality Data at Scale**
- - The pretraining corpus comprises large-scale, concept-balanced, diverse, and high-quality captions curated with strict filtering and quality control.
- - The instruction-tuning dataset is comprehensive and covers a wide range of tasks.
+### 3. Install PyTorch (CUDA 12.1)
-#### **Ultra-Efficient Training Framework**
- - The end-to-end training cost is about $16,000 on A100 GPUs at roughly $0.60 per GPU-hour.
- - The system is built on Megatron-LM with support for MoE, FP8, and long-sequence parallelism, and the codebase is optimized for cost-effective scaling.
+```bash
+pip install torch==2.3.0 torchvision==0.18.0 torchaudio==2.3.0 \
+ --index-url https://download.pytorch.org/whl/cu121
+```
-#### **Fully Open Framework**
- - The project releases high-quality pretraining and SFT datasets along with the complete training framework, configurations, and recipes.
- - It also provides detailed training logs and metrics to enable reproducibility and community adoption.
+### 4. Install core dependencies
+```bash
+pip install \
+ transformers==4.47.1 \
+ tokenizers \
+ sentencepiece \
+ einops \
+ timm \
+ Pillow \
+ numpy \
+ scipy \
+ tqdm \
+ pyyaml \
+ regex \
+ ftfy \
+ webdataset \
+ braceexpand \
+ protobuf "protobuf>=4.25.1,<7" \
+ wandb \
+ tensorboard \
+ open_clip_torch
+```
-## Models
+### 5. Install Transformer Engine (NVIDIA TE 2.11.0)
-| Model | HF Link | Training Log |
-|--------------------------|--------------------------------------------------------------------------------------------------------|-------------|
-| LLaVA-OneVision-1.5-4B-Instruct | [๐ค HF / 4B-Instruct](https://huggingface.co/lmms-lab/LLaVA-OneVision-1.5-4B-Instruct) | [๐ TensorBoard](https://huggingface.co/lmms-lab/LLaVA-OneVision-1.5-4B-Instruct/tensorboard) |
-| LLaVA-OneVision-1.5-8B-Instruct | [๐ค HF / 8B-Instruct](https://huggingface.co/lmms-lab/LLaVA-OneVision-1.5-8B-Instruct) | [๐ TensorBoard](https://huggingface.co/lmms-lab/LLaVA-OneVision-1.5-8B-Instruct/tensorboard) |
-| LLaVA-OneVision-1.5-4B-Base | [๐ค HF / 4B-Base](https://huggingface.co/lmms-lab/LLaVA-OneVision-1.5-4B-Base) | [๐ TensorBoard](https://huggingface.co/lmms-lab/LLaVA-OneVision-1.5-4B-Instruct/tensorboard) |
-| LLaVA-OneVision-1.5-8B-Base | [๐ค HF / 8B-Base](https://huggingface.co/lmms-lab/LLaVA-OneVision-1.5-8B-Base) | [๐ TensorBoard](https://huggingface.co/lmms-lab/LLaVA-OneVision-1.5-8B-Instruct/tensorboard) |
-## Datasets
+TE is required for fused attention and mixed-precision training:
-
-
- (a) The vocabulary coverage proportion in the LLaVA-OneVision-1.5 Mid-Training dataset before and after concept balancing.
- (b) Distribution of data sources within the LLaVA-OneVision-1.5 Mid-Training dataset.
- (c) Distribution of data sources within the LLaVA-OneVision-1.5 Instruct dataset.
-
+```bash
+# Set CUDA/cuDNN paths (adjust to your environment)
+export CUDA_HOME=/usr/local/cuda
+export CPATH="$CONDA_PREFIX/lib/python3.10/site-packages/nvidia/cudnn/include:$CPATH"
+export LIBRARY_PATH="$CONDA_PREFIX/lib/python3.10/site-packages/nvidia/cudnn/lib:$LIBRARY_PATH"
+export LD_LIBRARY_PATH="$CONDA_PREFIX/lib/python3.10/site-packages/nvidia/cudnn/lib:$LD_LIBRARY_PATH"
-| Description | Link | Status |
-|--------------------|--------------------------------------------------------------------------------------------------------|-------------|
-| LLaVA-OneVision-1.5-Mid-Training-85M | [๐คHF / Mid-Training 85M](https://huggingface.co/datasets/mvp-lab/LLaVA-OneVision-1.5-Mid-Training-85M) | Available |
-| LLaVA-OneVision-1.5-Instruct | [๐คHF / Instruct-Data](https://huggingface.co/datasets/mvp-lab/LLaVA-OneVision-1.5-Instruct-Data) | Available |
+pip install transformer_engine[pytorch]==2.11.0
+```
+
+### 6. Install Apex (fused CUDA kernels)
+```bash
+git clone https://github.com/NVIDIA/apex.git /tmp/apex
+cd /tmp/apex
+# The apex/ directory in this repo already contains the PyTorch 2.x patches
+# Copy patched files first
+cp /path/to/LLaVA-OneVision-1.5/apex/csrc/mlp.cpp csrc/
+cp /path/to/LLaVA-OneVision-1.5/apex/csrc/fused_dense.cpp csrc/
+
+pip install -v --disable-pip-version-check --no-cache-dir \
+ --no-build-isolation \
+ --config-settings "--build-option=--cpp_ext" \
+ --config-settings "--build-option=--cuda_ext" \
+ .
+```
-## Evaluation Results
+### 7. Install the AIAK Megatron submodule
+```bash
+cd /path/to/LLaVA-OneVision-1.5
+pip install -e aiak_megatron/
+```
-All evaluations were conducted using [lmms_eval](https://github.com/EvolvingLMMs-Lab/lmms-eval).
+### 8. Set PYTHONPATH
-
+```bash
+export PYTHONPATH=/path/to/LLaVA-OneVision-1.5:$PYTHONPATH
+```
+---
-## Quick Start with HuggingFace
+## Downloading Pretrained Checkpoints
-```python
-from transformers import AutoTokenizer, AutoProcessor, AutoModelForCausalLM
-from qwen_vl_utils import process_vision_info
-model_path = "lmms-lab/LLaVA-OneVision-1.5-8B-Instruct"
+### MobileLLM-R1-140M (Language Model)
-# default: Load the model on the available device(s)
-model = AutoModelForCausalLM.from_pretrained(
- model_path, torch_dtype="auto", device_map="auto", trust_remote_code=True
+```bash
+pip install huggingface_hub
+python -c "
+from huggingface_hub import snapshot_download
+snapshot_download(
+ repo_id='facebook/MobileLLM-R1-140M',
+ local_dir='checkpoints/MobileLLM-R1-140M'
)
+"
+```
-# default processor
-processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
-
-messages = [
- {
- "role": "user",
- "content": [
- {
- "type": "image",
- "image": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg",
- },
- {"type": "text", "text": "Describe this image."},
- ],
- }
-]
-
-# Preparation for inference
-text = processor.apply_chat_template(
- messages, tokenize=False, add_generation_prompt=True
-)
-image_inputs, video_inputs = process_vision_info(messages)
-inputs = processor(
- text=[text],
- images=image_inputs,
- videos=video_inputs,
- padding=True,
- return_tensors="pt",
-)
-inputs = inputs.to("cuda")
-
-# Inference: Generation of the output
-generated_ids = model.generate(**inputs, max_new_tokens=1024)
-generated_ids_trimmed = [
- out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
-]
-output_text = processor.batch_decode(
- generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
-)
-print(output_text)
+### FastViT MobileCLIP-L (Vision Encoder)
-```
+The FastViT encoder weights are part of Apple's ml-fastvlm release:
-## Evaluation
+```bash
+# Download from Apple ml-fastvlm (FastViT-HD 1.5B stage-3 checkpoint)
+# See: https://github.com/apple/ml-fastvlm
+# The vision tower weights are stored as mobileclip_l_1024 in the merged checkpoint below.
```
-# pip install git+https://github.com/EvolvingLMMs-Lab/lmms-eval.git
-accelerate launch --num_processes=8 --main_process_port 12399 -m lmms_eval \
- --model=llava_onevision1_5 \
- --model_args=pretrained=lmms-lab/LLaVA-OneVision-1.5-8B-Instruct,attn_implementation=flash_attention_2,max_pixels=3240000 \
- --tasks=mmmu_val,mmmu_pro_standard,mmbench_en_test,mmerealworld,mmerealworld_cn,ai2d,ai2d_no_mask,vstar_bench,chartqa,charxiv,docvqa_test,mathvista_testmini,mmstar,scienceqa \
- --batch_size=1
-```
+### Merged Checkpoint (MobileLLM + FastViT, Megatron format)
+The `mobilellm-fastvit-merged-tp1-pp1` checkpoint combines MobileLLM-R1-140M and FastViT
+into Megatron core format (TP=1, PP=1). Generate it once after downloading both models:
-## Quick Start Guide
+```bash
+# Convert MobileLLM from HuggingFace to Megatron format
+AIAK_TRAINING_PATH=$(pwd) python aiak_training_llm/models/mobilellm/megatron_checkpoint/convert_hf_to_mcore.py \
+ --hf-checkpoint checkpoints/MobileLLM-R1-140M \
+ --output checkpoints/mobilellm-fastvit-merged-tp1-pp1 \
+ --tp 1 --pp 1
+
+# The vision encoder weights are loaded automatically from the FastViT pretrained
+# weights specified by --vision-tower-name mobileclip_l_1024
+```
-### 1.๐ณ Docker (Recommended)
+---
-We strongly recommend using the docker environment for a seamless experience. The following instructions are tailored for the A100 80GB GPU environment.
+## Downloading the Dataset
+Stage 1 alignment uses the **LLaVA-558K** dataset in WebDataset format:
```bash
-# Clone repository
-git clone https://github.com/EvolvingLMMs-Lab/LLaVA-OneVision-1.5.git
-cd LLaVA-OneVision-1.5
+# Using HuggingFace hub
+python -c "
+from huggingface_hub import snapshot_download
+snapshot_download(
+ repo_id='lmms-lab/LLaVA-558K-Webdataset',
+ repo_type='dataset',
+ local_dir='data/LLaVA-558K-Webdataset'
+)
+"
+```
+
+The dataset is ~30 GB. It contains 558K image-text pairs in `.tar` WebDataset shards.
+
+---
+
+## Running Stage 1 Alignment Training
-docker build -t llava_megatron:25.04 .
+### Quick Start (single machine, 2 GPUs)
-# Run container with -w to set working directory directly to the mounted volume
-docker run -it --gpus all \
- --ipc host --net host --privileged --cap-add IPC_LOCK \
- --ulimit memlock=-1 --ulimit stack=67108864 --rm \
- -v $(pwd):/workspace/LLaVA-OneVision-1.5 \
- -w /workspace/LLaVA-OneVision-1.5 \
- --name "llava_megatron_container" \
- llava_megatron:25.04 /bin/bash
+```bash
+cd /path/to/LLaVA-OneVision-1.5
+
+# Set environment and run
+GPUS_PER_NODE=2 \
+DATA_PATH=data/LLaVA-558K-Webdataset \
+TOKENIZER_PATH=facebook/MobileLLM-R1-140M \
+PRETRAINED_CHECKPOINT=checkpoints/mobilellm-fastvit-merged-tp1-pp1 \
+bash examples/llava_ov_1_5/quick_start/stage_1_alignment_mobilellm_140m.sh
```
-### 2. Checkpoint and Format Conversion
+> **Note:** The script uses `torchrun`. Make sure your conda environment is active so `torchrun` is on your PATH.
+> If you get `torchrun: command not found`, prepend:
+> `PATH=/path/to/conda/envs/llava-mobile/bin:$PATH bash examples/...`
-You have two options to get started with LLaVA-OneVision-1.5-stage-0:
+### Script parameters
-#### Option 1: Download pre-trained model from Hugging Face
-Download our `LLaVA-OneVision-1.5-4B-stage0` model directly from [Hugging Face](https://huggingface.co/lmms-lab/LLaVA-OneVision-1.5-4B-stage0).
+The script accepts positional arguments:
-#### Option 2: Merge initial weights yourself
-Alternatively, you can merge the initial weights from the original ViT and LLM:
```bash
-python ds/merge_model.py \
---vit_path DeepGlint-AI/rice-vit-large-patch14-560 \
---llm_path Qwen/Qwen3-4B-Instruct-2507 \
---output LLaVA-OneVision-1.5-4B-stage0
+bash stage_1_alignment_mobilellm_140m.sh \
+ [TP=1] [PP=1] [SEQ_LEN=32768] [MBS=1] [GBS=2] [NSTEPS=500]
```
-Note: When merging weights, the adapter component will be initialized with default values.
-Convert the model from Hugging Face format to Megatron format:
+| Argument | Default | Description |
+|---|---|---|
+| `TP` | 1 | Tensor parallel degree |
+| `PP` | 1 | Pipeline parallel degree |
+| `SEQ_LEN` | 32768 | Maximum sequence length |
+| `MBS` | 1 | Micro batch size per GPU |
+| `GBS` | 2 | Global batch size |
+| `NSTEPS` | 1 | Number of training iterations |
+
+### Environment variable overrides
+
+| Variable | Default | Description |
+|---|---|---|
+| `GPUS_PER_NODE` | 2 | GPUs per machine |
+| `DATA_PATH` | `data/LLaVA-558K-Webdataset` | Dataset path |
+| `TOKENIZER_PATH` | `facebook/MobileLLM-R1-140M` | HF tokenizer |
+| `PRETRAINED_CHECKPOINT` | `checkpoints/mobilellm-fastvit-merged-tp1-pp1` | Megatron checkpoint |
+| `WANDB_API_KEY` | *(unset)* | Set to enable W&B logging |
+
+### What gets trained
+
+Stage 1 trains **only the 2-layer MLP adapter** (~3.5M parameters). Both the FastViT encoder and MobileLLM-R1-140M are fully frozen. This teaches the adapter to project visual features into the language model's embedding space.
+
+---
+
+## Demo: End-to-End Inference
+
+`inference_fastvlm.py` runs the full pipeline (FastViT โ Adapter โ MobileLLM) for a single image + text prompt without distributed training setup.
+
+### Usage
```bash
-AIAK_TRAINING_PATH=/workspace/LLaVA-OneVision-1.5 bash examples/llava_ov_1_5/convert/convert_4b_hf_to_mcore.sh \
-LLaVA-OneVision-1.5-4B-stage0 \
-LLaVA-OneVision-1.5-4B-stage0_mcore_tp1_pp1 \
-1 1
+CUDA_VISIBLE_DEVICES=0 python inference_fastvlm.py \
+ --image /path/to/image.jpg \
+ --prompt "<|image_pad|>\nDescribe what you see in this image." \
+ --checkpoint stage_1_alignment_mobilellm_140m/iter_0000500 \
+ --max-new-tokens 150
```
-### 3. Stage 1 Alignment-Training
+### Sample output (step-500 checkpoint, random noise image)
-Download LLaVA from [LLaVA-558K-Webdataset](https://huggingface.co/datasets/lmms-lab/LLaVA-558K-Webdataset).
+```
+Prompt: <|image_pad|>
+Describe what you see in this image.
+
+Generated: The image shows a [...]
+```
+### Without an image (text-only)
```bash
-# ============================================================
-# Required environment variables:
-# AIAK_TRAINING_PATH Root directory of the AIAK-Training-LLM project
-# DATA_PATH Directory with WebDataset shards (.tar) for pretraining
-# TOKENIZER_PATH Hugging Face tokenizer directory
-# CHECKPOINT_PATH Megatron-formatted checkpoint directory (e.g., mcore TP1/PP1)
-# SAVE_CKPT_PATH Output directory for saving training checkpoints
-AIAK_TRAINING_PATH=/workspace/LLaVA-OneVision-1.5 \
-DATA_PATH=LLaVA-558K-Webdataset \
-TOKENIZER_PATH=LLaVA-OneVision-1.5-4B-stage0 \
-CHECKPOINT_PATH=LLaVA-OneVision-1.5-4B-stage0_mcore_tp1_pp1 \
-bash examples/llava_ov_1_5/quick_start/stage_1_alignment_llava_ov_4b.sh
+CUDA_VISIBLE_DEVICES=0 python inference_fastvlm.py \
+ --prompt "What is the capital of France?" \
+ --max-new-tokens 50
```
-### 4. Stage 1.5 Mid-Training
+---
-Download our lightweight packed subset from [LLaVA-OneVision-1.5-Mid-Training-Quick-Start-3M-Webdataset](https://huggingface.co/datasets/lmms-lab/LLaVA-OneVision-1.5-Mid-Training-Webdataset-Quick-Start-3M).
+## Training Logs & Results
-```bash
-# ============================================================
-# Convert model to release format
-bash examples/llava_ov_1_5/convert/convert_4b_mcore_to_release.sh \
-stage_1_alignment_llava_ov_4b/iter_0002500/ \
-stage_1_alignment_llava_ov_4b_release 1 1
-# ============================================================
-# Launch
-AIAK_TRAINING_PATH=/workspace/LLaVA-OneVision-1.5 \
-DATA_PATH=LLaVA-OneVision-1.5-Mid-Training-Webdataset-Quick-Start-3M \
-TOKENIZER_PATH=LLaVA-OneVision-1.5-4B-stage0 \
-CHECKPOINT_PATH=stage_1_alignment_llava_ov_4b_release \
-bash examples/llava_ov_1_5/quick_start/stage_1.5_mid_training_llava_ov_4b.sh
+All training logs are stored in `stage_1_alignment_mobilellm_140m/`. The final 500-step run is:
+
+```
+stage_1_alignment_mobilellm_140m/run_2026-04-29_13:33:58_tp1_pp1_seqlen32768_mbs1_gbs4_500steps.log
```
+### Configuration for the 500-step run
-### 5. Stage 2 Instruct-Training
+| Setting | Value |
+|---|---|
+| GPUs | 4ร (TP=1, PP=1, DP=4) |
+| Global batch size | 4 samples |
+| Micro batch size | 1 sample/GPU |
+| Sequence length | 32,768 tokens |
+| Learning rate | 1e-4 (cosine decay to ~1e-5) |
+| Optimizer | Adam (ฮฒโ=0.9, ฮฒโ=0.99) |
+| Precision | BFloat16 |
+| Gradient clipping | 1.0 |
-Download LLaVA-NeXT-780k-webdataset at [LLaVA-NeXT-780K Dataset](https://huggingface.co/datasets/lmms-lab/LLaVA-NeXT-780k-webdataset).
+### Loss curve summary
-```bash
-# ============================================================
-# Convert model to release format
-bash examples/llava_ov_1_5/convert/convert_4b_mcore_to_release.sh \
-stage_1.5_mid_training_llava_ov_4b/iter_0020000/ \
-stage_1.5_mid_training_llava_ov_4b_release 1 1
-# ============================================================
-# # Launch
-AIAK_TRAINING_PATH=/workspace/LLaVA-OneVision-1.5 \
-DATA_PATH=LLaVA-NeXT-780k-Webdataset \
-TOKENIZER_PATH=LLaVA-OneVision-1.5-4B-stage0 \
-CHECKPOINT_PATH=stage_1.5_mid_training_llava_ov_4b_release \
-bash examples/llava_ov_1_5/quick_start/stage_2_instruct_llava_ov_4b.sh
-```
-
-
-### 6. Convert mcore to Hugging Face
-```bash
-AIAK_TRAINING_PATH=/workspace/LLaVA-OneVision-1.5 \
-bash examples/llava_ov_1_5/convert/convert_4b_mcore_to_hf.sh \
-stage_2_instruct_llava_ov_4b/iter_0003500 \
-LLaVA-OneVision-1.5-4B-3M-Mid-Training-780K-Instruct \
-1 1
-# Copy non-model files (e.g., tokenizer config) to the new directory
-find LLaVA-OneVision-1.5-4B-stage0/ -type f -not -iname '*safetensors*' -exec cp {} LLaVA-OneVision-1.5-4B-3M-Mid-Training-780K-Instruct/ ';'
+| Iteration | LM Loss |
+|---|---|
+| 1 | ~11.78 (random adapter init) |
+| 100 | ~11.2x |
+| 500 | **11.30** |
+
+> The loss decrease is expected to be gradual at Stage 1 since only the 3.5M-parameter adapter is trained and the LLM (140M params) is frozen. Full convergence requires Stage 2 instruction fine-tuning.
+
+### Checkpoint
+
+The latest checkpoint is at iteration 500 (Megatron format, TP=1 PP=1).
+The binary weights are **not stored in git** (too large; 529 MB + 129 MB).
+They are available on the training server at:
+
+```
+/share/data/drive_3/mobile_vlm/LLaVA-OneVision-1.5/stage_1_alignment_mobilellm_140m/iter_0000500/
+โโโ mp_rank_00/model_optim_rng.pt # full model + optimizer state (529 MB)
+โโโ mp_rank_00/distrib_optim.pt # distributed optimizer shards (129 MB)
```
-### 7. Evaluation
+To resume training from this checkpoint, set:
```bash
-# pip install git+https://github.com/EvolvingLMMs-Lab/lmms-eval.git
-CUDA_VISIBLE_DEVICES=4,5,6,7 accelerate launch \
---num_processes=4 --main_process_port 12399 -m lmms_eval --model=llava_onevision1_5 --batch_size=1 --tasks=mme \
---model_args=pretrained=/workspace/LLaVA-OneVision-1.5/LLaVA-OneVision-1.5-4B-3M-Mid-Training-780K-Instruct,max_pixels=3240000
-```
-
-## Fully Reproducing Guide
-
-> [!TIP]
-> More detailed reproduction steps for the complete process will be provided after the dataset upload is completed.
-
-
-### Mid-Training
-
-To improve model training efficiency, we implement offline sample packing:
-
-1. Download the [**Mid-Training-85M Dataset**](https://huggingface.co/datasets/lmms-lab/LLaVA-One-Vision-1.5-Mid-Training-85M)
-2. Pack the data into WebDataset format, refer to [**Examples offlinepacking**](examples_offline_packing) and [**Offline Padding-Free Data Packing**](examples/llava_ov_1_5/sample_packing/README.md)
-
-
-### Instruct
-1. Download the [**LLaVA-OneVision-1.5-Instruct-Data**](https://huggingface.co/datasets/lmms-lab/LLaVA-OneVision-1.5-Instruct-Data)
-2. Convert the data into WebDataset format, refer to [**Conversion for Mixed Instruction Data**](docs/sft_data_preprocessing.md)
-
-## Roadmap
-
-Q4 2025 Key Deliverables:
-
-1. **Ultra-efficient MoE Training**
-2. **Full Video Input LLM**
-
-
-## Contributors
-Thanks so much to all of our amazing contributors!
-
-
-
-
-
-## Citation
-
-If you find *LLaVA-OneVision-1.5* useful in your research, please consider to cite the following related papers:
-
+PRETRAINED_CHECKPOINT=stage_1_alignment_mobilellm_140m/iter_0000500
```
+
+---
+
+## Modifications vs. Upstream
+
+This branch modifies the original LLaVA-OneVision-1.5 codebase in the following ways. All changes are clearly separated from the upstream by the file paths below.
+
+### New files added
+
+| File | Description |
+|---|---|
+| `aiak_training_llm/models/mobilellm/mobilellm_model.py` | MobileLLM-R1-140M Megatron wrapper |
+| `aiak_training_llm/models/mobilellm/mobilellm_config.py` | Reads HF `config.json` โ `TransformerConfig`; documents NoPEโRoPE divergence |
+| `aiak_training_llm/models/mobilellm/mobilellm_layer_spec.py` | Transformer layer spec with correct QK-norm (checks `config.qk_layernorm`) |
+| `aiak_training_llm/models/mobilellm/mobilellm_provider.py` | Megatron model provider |
+| `aiak_training_llm/models/fastvit/fastvit_vision_model.py` | Megatron-compatible FastViT wrapper |
+| `aiak_training_llm/models/fastvit/mobileclip_encoder.py` | `MobileCLIPVisionTower` that loads Apple FastViT weights |
+| `aiak_training_llm/models/fastvit/fastvit_preprocessor.py` | Image preprocessing for FastViT (1024 px, square padding) |
+| `aiak_training_llm/models/fastvit/mobileclip/` | Apple ml-fastvlm FastViT model code |
+| `inference_fastvlm.py` | Stand-alone inference demo |
+| `examples/llava_ov_1_5/quick_start/stage_1_alignment_mobilellm_140m.sh` | Training launcher for MobileLLM + FastViT |
+
+### Modified files (vs. upstream)
+
+| File | Change |
+|---|---|
+| `aiak_training_llm/models/llavaov_1_5/llavaov_1_5_model.py` | Added 6-step pipeline logging; replaced `[BIG DEBUG]` prints with structured `[1/6]โฆ[6/6]` output; added per-step token accuracy |
+| `aiak_training_llm/models/llavaov_1_5/llavaov_1_5_layer_spec.py` | Removed duplicate `get_mobilellm_layer_with_te_spec` that hard-coded `q_layernorm=IdentityOp` (bug: ignored config); now imports the correct version from `mobilellm_layer_spec.py` |
+| `aiak_training_llm/models/llavaov_1_5/llavaov_1_5_config.py` | Sets `qk_layernorm=True` for MobileLLM (matching `use_qk_norm: true` in HF config) |
+| `aiak_training_llm/data/multimodal/qwen2vl_task_encoder.py` | Added FastViT preprocessing path (`--use-fastvit` flag); per-image debug prints removed |
+| `aiak_megatron/megatron/core/transformer/dot_product_attention.py` | Fixed GQA shape handling for non-square query/key heads; removed debug prints |
+| `aiak_megatron/megatron/core/transformer/attention.py` | Removed debug prints |
+| `aiak_megatron/megatron/core/extensions/transformer_engine.py` | ROCm/AMD compatibility: skip CUDA arch checks on AMD GPUs |
+| `aiak_megatron/megatron/core/extensions/transformer_engine2.py` | ROCm/AMD compatibility |
+| `apex/csrc/mlp.cpp` | PyTorch 2.x API fix: `.type()` โ `.scalar_type()`, `.options()` |
+| `apex/csrc/fused_dense.cpp` | PyTorch 2.x API fix: same as above |
+
+### Deleted dead code files
+
+| File | Reason |
+|---|---|
+| `aiak_training_llm/models/llavaov_1_5/adapter.py` | Never imported; active adapter is `qwen_vl/adapter.py` |
+| `aiak_training_llm/models/llavaov_1_5/llavaov_1_5_layer_spec_org.py` | Old backup copy |
+| `aiak_training_llm/models/llavaov_1_5/llavaov_mobilellm_config.py` | Superseded by `mobilellm/mobilellm_config.py` |
+
+### Key architectural decisions
+
+1. **QK LayerNorm (bug fix):** MobileLLM-R1-140M uses `use_qk_norm: true` in its HuggingFace config. The upstream code ignored this, causing attention to run without QK normalization. Fixed in `llavaov_1_5_layer_spec.py`.
+
+2. **NoPE โ RoPE (intentional divergence):** MobileLLM-R1 is a NoPE (No Positional Encoding) model โ all 15 layers have `no_rope_layers=[1,1,...,1]`. This Megatron integration applies RoPE with `rope_theta=8000000` (from the HF config) to provide positional encoding for multimodal sequences. This is documented in `mobilellm_config.py`.
+
+3. **FastViT output shape:** FastViT MobileCLIP-L outputs 4D features `[B, 3072, H, W]` where `H=W=16` (for 1024 px input with patch_size=64). These are reshaped to `[B, 256, 3072]` before the adapter.
+
+4. **Vision token count mismatch:** The data pipeliner packs multiple images per sequence. A trim/pad mechanism in `llavaov_1_5_model.py` ensures the adapter output token count always matches the number of `<|image_pad|>` tokens in `input_ids`.
+
+---
+
+## Credits & Citations
+
+This project builds on top of:
+
+### LLaVA-OneVision-1.5 (base framework)
+
+```bibtex
@inproceedings{LLaVA-OneVision-1.5,
title={LLaVA-OneVision-1.5: Fully Open Framework for Democratized Multimodal Training},
- author={An, Xiang and Xie, Yin and Yang, Kaicheng and Zhang, Wenkang and Zhao, Xiuwei and Cheng, Zheng and Wang, Yirui and Xu, Songcen and Chen, Changrui and Wu, Chunsheng and Tan, Huajie and Li, Chunyuan and Yang, Jing and Yu, Jie and Wang, Xiyao and Qin, Bin and Wang, Yumeng and Yan, Zizhen and Feng, Ziyong and Liu, Ziwei and Li, Bo and Deng, Jiankang},
- booktitle={arXiv},
- year={2025}
- }
-
-@inproceedings{xie2025region,
- title={Region-based Cluster Discrimination for Visual Representation Learning},
- author={Xie, Yin and Yang, Kaicheng and An, Xiang and Wu, Kun and Zhao, Yongle and Deng, Weimo and Ran, Zimin and Wang, Yumeng and Feng, Ziyong and Miles, Roy and Elezi, Ismail and Deng, Jiankang},
- booktitle={ICCV},
+ author={An, Xiang and Xie, Yin and Yang, Kaicheng and others},
+ booktitle={arXiv},
year={2025}
}
+```
+
+GitHub: https://github.com/EvolvingLMMs-Lab/LLaVA-OneVision-1.5
-@article{lillava,
- title={LLaVA-OneVision: Easy Visual Task Transfer},
- author={Li, Bo and Zhang, Yuanhan and Guo, Dong and Zhang, Renrui and Li, Feng and Zhang, Hao and Zhang, Kaichen and Zhang, Peiyuan and Li, Yanwei and Liu, Ziwei and Li, Chunyuan},
- journal={Transactions on Machine Learning Research}
+### MobileLLM-R1-140M (Language Model)
+
+```bibtex
+@article{mobilellm,
+ title={MobileLLM: Optimizing Sub-billion Parameter Language Models for On-Device Use Cases},
+ author={Liu, Zechun and others},
+ journal={arXiv},
year={2024}
}
```
-## Acknowledgement
+HuggingFace: https://huggingface.co/facebook/MobileLLM-R1-140M
-We extend our sincere gratitude to **AIAK team of the** [**Baige AI computing platform**](https://cloud.baidu.com/product/aihc.html) **from Baidu AI Cloud** for providing the exceptional training framework. The outstanding capabilities of AIAK-Training-LLM and AIAK-Megatron have significantly accelerated our training process with remarkable efficiency. These cutting-edge frameworks have been instrumental in achieving our research goals. `To get full AIAK support, you can contact Baidu Cloud.`
+### FastVLM / MobileCLIP-L (Vision Encoder)
-We acknowledge the support of [Synvo AI](https://synvo.ai/) for contributing to the partial data annotation in this work, and also thank the maintainers and contributors of the following open-source projects, whose work greatly inspired and supported our research:
+```bibtex
+@article{fastvlm,
+ title={FastVLM: Efficient Vision Encoding for Vision Language Models},
+ author={Prabhu, Shirin and others},
+ journal={CVPR},
+ year={2025}
+}
+```
+
+GitHub: https://github.com/apple/ml-fastvlm
+
+### Megatron-LM (Training Framework)
+
+GitHub: https://github.com/NVIDIA/Megatron-LM
+
+---
-- LLaVA: Large Language-and-Vision Assistant โ [LLaVA](https://github.com/haotian-liu/LLaVA)
-- LLaVA-NeXT: Next-generation multi-modal assistant โ [LLaVA-NeXT](https://github.com/LLaVA-VL/LLaVA-NeXT)
-- lmms-eval: A standardized evaluation framework for Large Multimodal Models โ [lmms-eval](https://github.com/EvolvingLMMs-Lab/lmms-eval)
-- Megatron-LM: Efficient, scalable training for large language models โ [Megatron-LM](https://github.com/NVIDIA/Megatron-LM)
-- Qwen2.5-VL: Strong vision-language foundation model โ [Qwen2.5-VL](https://github.com/QwenLM/Qwen2.5-VL)
-- InternVL: Open-source large-scale vision-language foundation model โ [InternVL](https://github.com/OpenGVLab/InternVL)
-- Qwen3: Next-generation Qwen LLM โ [Qwen](https://github.com/QwenLM/Qwen)
-- MetaCLIP: Scalable contrastive pretraining โ [MetaCLIP](https://github.com/facebookresearch/MetaCLIP)
-- FineVision: Open Data Is All You Need โ [FineVision](https://huggingface.co/spaces/HuggingFaceM4/FineVision)
+*For questions or issues, please open a GitHub issue on this repository.*
diff --git a/Stage1/alignment.sh b/Stage1/alignment.sh
new file mode 100755
index 00000000..6f8fd25c
--- /dev/null
+++ b/Stage1/alignment.sh
@@ -0,0 +1,69 @@
+#!/bin/bash
+#SBATCH --job-name=llava_stage1_4b
+#SBATCH --time=72:00:00
+#SBATCH --nodes=1
+#SBATCH -p long
+#SBATCH -q gpu-12
+#SBATCH --gres=gpu:4 # 4 A100 GPUs on a single node
+#SBATCH --mem=230G # node RAM, not GPU RAM
+#SBATCH --ntasks-per-node=4 # one task per GPU
+#SBATCH --cpus-per-task=16 # adjust if you want fewer CPU cores
+#SBATCH --output=/l/users/rana.zayed/new_fastvlm/LLaVA-OneVision-1.5/Stage1/logs/%x-%j.out # logs/llava_stage1_4b-.out
+
+# ---- ENV SETUP ----
+source ~/miniconda3/etc/profile.d/conda.sh
+conda activate llava-ov-4b-clean
+# conda activate apex_cuda120
+export PATH=/usr/local/cuda-12.1/bin:$PATH
+export LD_LIBRARY_PATH=/usr/local/cuda-12.1/lib64:$LD_LIBRARY_PATH
+export APEX_CUDA_EXT=1
+
+# Go to repo root on CIAI cluster
+# cd /l/users/rana.zayed/new_fastvlm/LLaVA-OneVision-1.5
+# # Go to repo root on 156 machine
+cd /share/data/drive_3/mobile_vlm/LLaVA-OneVision-1.5
+
+# ============================================================
+# Required environment variables:
+# AIAK_TRAINING_PATH Root directory of the AIAK-Training-LLM project
+# DATA_PATH Directory with WebDataset shards (.tar) for pretraining
+# TOKENIZER_PATH Hugging Face tokenizer directory
+# CHECKPOINT_PATH Megatron-formatted checkpoint directory (e.g., mcore TP1/PP1)
+# SAVE_CKPT_PATH Output directory for saving training checkpoints
+# export CUDA_HOME=/apps/local/nvidia/cuda-12.0
+# export PATH="$CUDA_HOME/bin:$PATH"
+# export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH
+
+# Use CUDA 12.1 libraries but system nvcc (CUDA 10.1) since 12.1 doesn't have nvcc
+export CUDA_HOME=/usr
+export PATH="/usr/bin:$PATH"
+export LD_LIBRARY_PATH="/usr/local/cuda-12.1/lib64:/usr/local/cuda-12.1/targets/x86_64-linux/lib:$LD_LIBRARY_PATH"
+
+
+# AIAK_TRAINING_PATH=/l/users/rana.zayed/new_fastvlm/LLaVA-OneVision-1.5 \
+# DATA_PATH=/l/users/rana.zayed/new_fastvlm/LLaVA-OneVision-1.5/data/LLaVA-558K-Webdataset \
+# TOKENIZER_PATH=/l/users/rana.zayed/new_fastvlm/LLaVA-OneVision-1.5/checkpoints/LLaVA-OneVision-1.5-4B-stage0 \
+# CHECKPOINT_PATH=/l/users/rana.zayed/new_fastvlm/LLaVA-OneVision-1.5/checkpoints/LLaVA-OneVision-1.5-4B-stage0_mcore_tp1_pp1 \
+
+# echo "AIAK_TRAINING_PATH=${AIAK_TRAINING_PATH}"
+# echo "DATA_PATH=${DATA_PATH}"
+# echo "TOKENIZER_PATH=${TOKENIZER_PATH}"
+# echo "CHECKPOINT_PATH=${CHECKPOINT_PATH}"
+# echo "SLURM_NODELIST=${SLURM_NODELIST}"
+
+# Weights & Biases configuration
+export WANDB_API_KEY="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"
+export WANDB_PROJECT="llava-ov-1_5"
+export WANDB_NAME="mobilellm_integration"
+
+export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
+export GPUS_PER_NODE=8
+export MASTER_PORT=26000
+
+# Choose which backbone to use for training:
+# ============================================================
+# Option 1: MobileLLM-R1-140M (140M params, efficient for mobile/edge)
+bash examples/llava_ov_1_5/quick_start/stage_1_alignment_mobilellm_140m.sh
+
+# Option 2: Original Qwen2.5-4B backbone (4B params)
+# bash examples/llava_ov_1_5/quick_start/stage_1_alignment_llava_ov_4b.sh
\ No newline at end of file
diff --git a/Stage1/alignment_rocm.sh b/Stage1/alignment_rocm.sh
new file mode 100755
index 00000000..4657de3a
--- /dev/null
+++ b/Stage1/alignment_rocm.sh
@@ -0,0 +1,78 @@
+#!/bin/bash
+
+#SBATCH --job-name=llava_stage1_4b_amd
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=1
+#SBATCH --cpus-per-task=128
+#SBATCH --gres=gpu:8
+#SBATCH --qos=skqos
+#SBATCH --partition=faculty
+#SBATCH --output=/vast/users/salman.khan/mobile_vlm/llava_ov1.5/LLaVA-OneVision-1.5/Stage1/logs/%x-%j.out
+
+# ---- ENV SETUP (AMD) ----
+source ~/.bashrc
+conda activate mobile_vlm
+
+export MIOPEN_DISABLE_CACHE=1
+export PYTORCH_TUNABLEOP_ENABLED=0
+
+export ROCM_HOME=${ROCM_HOME:-/opt/rocm}
+export PATH="${ROCM_HOME}/bin:${PATH}"
+export LD_LIBRARY_PATH="${ROCM_HOME}/lib:${ROCM_HOME}/lib64:${LD_LIBRARY_PATH}"
+
+export HIP_VISIBLE_DEVICES=${HIP_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}
+
+# Force HuggingFace offline mode to use local files only
+#export HF_HUB_OFFLINE=1
+#export TRANSFORMERS_OFFLINE=1
+
+
+# RCCL/NCCL runtime hints (tune as needed)
+export NCCL_DEBUG=${NCCL_DEBUG:-WARN}
+export NCCL_COLLNET_ENABLE=${NCCL_COLLNET_ENABLE:-0}
+export NCCL_P2P_ENABLE=${NCCL_P2P_ENABLE:-1}
+# export NCCL_SOCKET_IFNAME=eno1 # uncomment and set to your NIC if needed
+
+# Resolve repo root relative to this script (Stage1/..)
+SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
+REPO_ROOT=/vast/users/salman.khan/mobile_vlm/llava_ov1.5/LLaVA-OneVision-1.5
+cd "$REPO_ROOT" || exit 1
+
+
+# Go to repo root
+cd "$REPO_ROOT" || { echo "[Error] Repo root not found: $REPO_ROOT"; exit 1; }
+
+echo "=== ENV CHECK ==="
+which conda
+which python
+python -V
+echo "CONDA_DEFAULT_ENV=$CONDA_DEFAULT_ENV"
+echo "CONDA_PREFIX=$CONDA_PREFIX"
+python -c "import sys; print('sys.executable=', sys.executable)"
+python -c "import torch; print('torch=', torch.__version__)" || echo "TORCH NOT FOUND"
+pip -V
+pip list | grep -E "torch|pytorch" || true
+echo "=== END ENV CHECK ==="
+
+# Required environment variables
+export AIAK_TRAINING_PATH="${AIAK_TRAINING_PATH:-$REPO_ROOT}"
+export AIAK_MAGATRON_PATH="${AIAK_MAGATRON_PATH:-$REPO_ROOT/aiak_megatron}"
+export DATA_PATH="${DATA_PATH:-$REPO_ROOT/data/LLaVA-558K-Webdataset}"
+# Use MobileLLM tokenizer and checkpoint (let stage_1_alignment_mobilellm_140m.sh set defaults)
+export TOKENIZER_PATH="${TOKENIZER_PATH:-facebook/MobileLLM-R1-140M}"
+export PRETRAINED_CHECKPOINT="${PRETRAINED_CHECKPOINT:-$REPO_ROOT/checkpoints/mobilellm-fastvit-merged-tp1-pp1}"
+# Add megatron to PYTHONPATH so imports work
+export PYTHONPATH="${AIAK_MAGATRON_PATH}:${AIAK_TRAINING_PATH}:${PYTHONPATH}"
+
+echo "AIAK_TRAINING_PATH=${AIAK_TRAINING_PATH}"
+echo "AIAK_MAGATRON_PATH=${AIAK_MAGATRON_PATH}"
+echo "DATA_PATH=${DATA_PATH}"
+echo "TOKENIZER_PATH=${TOKENIZER_PATH}"
+echo "SLURM_NODELIST=${SLURM_NODELIST}"
+echo "PYTHONPATH=${PYTHONPATH}"
+
+# Weights & Biases configuration
+export WANDB_API_KEY="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"
+export WANDB_PROJECT="llava-ov-1_5"
+export WANDB_NAME="fastvit_integration"
+bash examples/llava_ov_1_5/quick_start/stage_1_alignment_mobilellm_140m.sh
\ No newline at end of file
diff --git a/Stage1/inference_fastvlm.py b/Stage1/inference_fastvlm.py
new file mode 100644
index 00000000..f70d7d94
--- /dev/null
+++ b/Stage1/inference_fastvlm.py
@@ -0,0 +1,114 @@
+"""
+FastVLM Inference Script
+Example run:
+python inference_fastvlm.py --checkpoint_path ./stage_1_alignment_llava_ov_4b/iter_0000020 \
+ --image_path ./test_image.jpg \
+ --prompt "What is in this image?"
+"""
+
+import os
+import sys
+import torch
+from PIL import Image
+from argparse import ArgumentParser
+
+# Add repo root to path
+REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, REPO_ROOT)
+sys.path.insert(0, os.path.join(REPO_ROOT, 'aiak_megatron'))
+
+from transformers import AutoProcessor
+from aiak_training_llm.models.fastvit.fastvit_preprocessor import FastViTImageProcessor
+from aiak_training_llm.models.fastvit.mm_utils import expand2square
+
+# Argument parser
+parser = ArgumentParser(description="FastVLM Inference")
+parser.add_argument('--checkpoint_path', type=str,
+ default='/share/data/drive_3/mobile_vlm/LLaVA-OneVision-1.5/stage_1_alignment_llava_ov_4b/iter_0000020',
+ help='Path to trained checkpoint directory')
+parser.add_argument('--tokenizer_path', type=str,
+ default='/share/data/drive_3/mobile_vlm/LLaVA-OneVision-1.5/checkpoints/LLaVA-OneVision-1.5-4B-stage0',
+ help='Path to tokenizer')
+parser.add_argument('--image_path', type=str, default='test_image.jpg',
+ help='Path to input image')
+parser.add_argument('--prompt', type=str, default='What is in this image?',
+ help='Text prompt for the model')
+parser.add_argument('--image_size', type=int, default=1024,
+ help='FastViT image size (384 or 1024)')
+parser.add_argument('--use_gpu', action='store_true', default=True,
+ help='Use GPU for inference')
+args = parser.parse_args()
+
+print("=" * 80)
+print("FastVLM Inference")
+print("=" * 80)
+print(f"Checkpoint: {args.checkpoint_path}")
+print(f"Tokenizer: {args.tokenizer_path}")
+print(f"Image: {args.image_path}")
+print(f"Prompt: {args.prompt}")
+print(f"Image Size: {args.image_size}")
+print("=" * 80)
+
+# Device setup
+device = torch.device("cuda:0" if torch.cuda.is_available() and args.use_gpu else "cpu")
+print(f"Using device: {device}")
+
+# Load tokenizer/processor
+print("\nLoading tokenizer and processor...")
+processor = AutoProcessor.from_pretrained(args.tokenizer_path, trust_remote_code=True)
+tokenizer = processor.tokenizer
+
+# Initialize FastViT image processor
+fastvit_processor = FastViTImageProcessor(image_size=args.image_size)
+print(f"FastViT processor initialized with image_size={args.image_size}")
+
+# Load and preprocess image
+print(f"\nLoading image from: {args.image_path}")
+if not os.path.exists(args.image_path):
+ print(f"ERROR: Image file not found: {args.image_path}")
+ print("Please provide a valid image path using --image_path")
+ sys.exit(1)
+
+image = Image.open(args.image_path).convert('RGB')
+print(f"Original image size: {image.size}")
+
+# Preprocess with FastViT (pad to square)
+mean_color = tuple(int(x * 255) for x in fastvit_processor.image_mean)
+image_padded = expand2square(image, mean_color)
+print(f"Padded to square: {image_padded.size}")
+
+pixel_values = fastvit_processor(image_padded)
+print(f"Preprocessed image shape: {pixel_values.shape}")
+
+# Create prompt with vision tokens
+IMAGE_TOKEN = "<|image_pad|>"
+VISION_START = "<|vision_start|>"
+VISION_END = "<|vision_end|>"
+
+# Format: <|im_start|>system\n...<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|>...<|vision_end|>\nPROMPT<|im_end|>\n<|im_start|>assistant\n
+conversation = f"<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n{VISION_START}{IMAGE_TOKEN}{VISION_END}\n{args.prompt}<|im_end|>\n<|im_start|>assistant\n"
+
+print("\nTokenizing prompt...")
+input_ids = tokenizer(conversation, return_tensors="pt")["input_ids"]
+print(f"Input IDs shape: {input_ids.shape}")
+print(f"Prompt tokens: {input_ids.shape[1]}")
+
+# TODO: Load your trained model checkpoint here
+# This requires implementing model loading from Megatron checkpoint
+print("\n" + "=" * 80)
+print("NOTE: Model loading from Megatron checkpoint not yet implemented.")
+print("This script currently only demonstrates preprocessing.")
+print("\nTo complete inference, you need to:")
+print("1. Load the model from checkpoint using Megatron utilities")
+print("2. Convert distributed checkpoint to single GPU format")
+print("3. Call model.forward() with preprocessed inputs")
+print("=" * 80)
+
+# Placeholder for model inference
+print("\nPreprocessed inputs ready:")
+print(f" - pixel_values: {pixel_values.shape} ({pixel_values.dtype})")
+print(f" - input_ids: {input_ids.shape}")
+print(f" - Device: {device}")
+
+print("\nInference complete (preprocessing only).")
+
\ No newline at end of file
diff --git a/Stage1/test_image.jpg b/Stage1/test_image.jpg
new file mode 100644
index 00000000..02a0f92f
Binary files /dev/null and b/Stage1/test_image.jpg differ
diff --git a/Stage1/test_inference_mobilellm.py b/Stage1/test_inference_mobilellm.py
new file mode 100644
index 00000000..940e4311
--- /dev/null
+++ b/Stage1/test_inference_mobilellm.py
@@ -0,0 +1,211 @@
+"""
+Test FastVLM inference with MobileLLM-140M + FastViT
+Uses the same model setup as training but loads from checkpoint for inference.
+"""
+import os
+import sys
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from argparse import ArgumentParser
+
+# Add paths
+REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, REPO_ROOT)
+sys.path.insert(0, os.path.join(REPO_ROOT, 'aiak_megatron'))
+sys.path.insert(0, os.path.join(REPO_ROOT, 'aiak_training_llm'))
+
+from transformers import AutoTokenizer
+from aiak_training_llm.models.fastvit.fastvit_preprocessor import FastViTImageProcessor
+from aiak_training_llm.models.fastvit.mm_utils import expand2square
+from aiak_training_llm.models.llavaov_1_5.llavaov_1_5_config import llava_ov_mobilellm_140m
+from aiak_training_llm.models.llavaov_1_5.llavaov_1_5_model import LlavaOnevision1_5
+from aiak_training_llm.models.mobilellm.mobilellm_layer_spec import get_mobilellm_layer_with_te_spec
+from aiak_training_llm.models.qwen.layer_spec import get_adapeter_layer_with_spec
+
+def main():
+ parser = ArgumentParser(description="FastVLM + MobileLLM Inference Test")
+ parser.add_argument('--checkpoint', type=str,
+ default='checkpoints/mobilellm-fastvit-merged-tp1-pp1',
+ help='Path to merged checkpoint')
+ parser.add_argument('--tokenizer', type=str,
+ default='checkpoints/MobileLLM-R1-140M',
+ help='Path to tokenizer')
+ parser.add_argument('--image', type=str, default='test_image.jpg',
+ help='Path to test image')
+ parser.add_argument('--prompt', type=str, default='What is in this image?',
+ help='Text prompt')
+ parser.add_argument('--image_size', type=int, default=1024,
+ help='FastViT image size')
+ parser.add_argument('--max_new_tokens', type=int, default=100,
+ help='Maximum number of tokens to generate')
+ parser.add_argument('--temperature', type=float, default=0.7,
+ help='Sampling temperature')
+ args = parser.parse_args()
+
+ print("=" * 80)
+ print("FastVLM + MobileLLM-140M Inference Test")
+ print("=" * 80)
+ print(f"Checkpoint: {args.checkpoint}")
+ print(f"Tokenizer: {args.tokenizer}")
+ print(f"Image: {args.image}")
+ print(f"Image Size: {args.image_size}")
+ print("=" * 80)
+
+ # Check checkpoint exists
+ checkpoint_path = os.path.join(REPO_ROOT, args.checkpoint)
+ if not os.path.exists(checkpoint_path):
+ print(f"\nERROR: Checkpoint not found: {checkpoint_path}")
+ print("Please run training first or provide a valid checkpoint path.")
+ return
+
+ # Load tokenizer
+ print("\n[1/4] Loading tokenizer...")
+ tokenizer_path = os.path.join(REPO_ROOT, args.tokenizer)
+ tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, trust_remote_code=True)
+ print(f"โ Loaded tokenizer from {tokenizer_path}")
+ print(f" Vocab size: {len(tokenizer)}")
+
+ # Load and preprocess image
+ print("\n[2/4] Loading and preprocessing image...")
+ image_path = os.path.join(os.path.dirname(__file__), args.image) if not os.path.isabs(args.image) else args.image
+ if not os.path.exists(image_path):
+ print(f"ERROR: Image not found: {image_path}")
+ return
+
+ image = Image.open(image_path).convert('RGB')
+ print(f" Original size: {image.size}")
+
+ # Initialize FastViT processor
+ fastvit_processor = FastViTImageProcessor(image_size=args.image_size)
+ mean_color = tuple(int(x * 255) for x in fastvit_processor.image_mean)
+ image_padded = expand2square(image, mean_color)
+ pixel_values = fastvit_processor(image_padded).unsqueeze(0) # Add batch dim
+ print(f" Padded to square: {image_padded.size}")
+ print(f" Preprocessed shape: {pixel_values.shape}")
+
+ # Create prompt
+ print("\n[3/4] Tokenizing prompt...")
+ IMAGE_TOKEN = "<|image_pad|>"
+ VISION_START = "<|vision_start|>"
+ VISION_END = "<|vision_end|>"
+
+ conversation = (
+ "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
+ f"<|im_start|>user\n{VISION_START}{IMAGE_TOKEN}{VISION_END}\n{args.prompt}<|im_end|>\n"
+ "<|im_start|>assistant\n"
+ )
+
+ input_ids = tokenizer(conversation, return_tensors="pt")["input_ids"]
+ print(f" Prompt: {args.prompt}")
+ print(f" Input IDs shape: {input_ids.shape}")
+
+ # Initialize model
+ print("\n[4/4] Loading model and running inference...")
+ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+ print(f" Using device: {device}")
+
+ # Get model configs
+ print(" Building model...")
+ language_config, vision_config, adapter_config = llava_ov_mobilellm_140m()
+
+ # Get layer specs (FastViT doesn't use layer spec, uses nn.Module directly)
+ language_layer_spec = get_mobilellm_layer_with_te_spec()
+ vision_layer_spec = None # FastViT uses direct nn.Module
+ adapter_layer_spec = get_adapeter_layer_with_spec()
+
+ # Build model
+ model = LlavaOnevision1_5(
+ language_config=language_config,
+ vision_config=vision_config,
+ adapter_config=adapter_config,
+ language_layer_spec=language_layer_spec,
+ vision_layer_spec=vision_layer_spec,
+ adapter_layer_spec=adapter_layer_spec,
+ language_vocab_size=128256,
+ language_max_sequence_length=512,
+ pre_process=True,
+ post_process=True,
+ fp16_lm_cross_entropy=False,
+ parallel_output=False,
+ share_embeddings_and_output_weights=True
+ )
+
+ # Load checkpoint
+ print(f" Loading weights from {checkpoint_path}...")
+ checkpoint_file = os.path.join(checkpoint_path, 'release', 'mp_rank_00', 'model_optim_rng.pt')
+ if os.path.exists(checkpoint_file):
+ checkpoint = torch.load(checkpoint_file, map_location='cpu', weights_only=False)
+ # Remove 'module.' prefix if present
+ state_dict = {}
+ for k, v in checkpoint['model'].items():
+ new_key = k.replace('module.', '') if k.startswith('module.') else k
+ state_dict[new_key] = v
+
+ missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
+ print(f" โ Checkpoint loaded")
+ print(f" Missing keys: {len(missing_keys)} (adapter weights expected)")
+ print(f" Unexpected keys: {len(unexpected_keys)}")
+ else:
+ print(f" WARNING: Checkpoint file not found: {checkpoint_file}")
+ print(" Using randomly initialized weights")
+
+ model = model.to(device)
+ model.eval()
+ print(" โ Model ready")
+
+ # Move inputs to device
+ pixel_values = pixel_values.to(device)
+ input_ids = input_ids.to(device)
+
+ # Generate
+ print("\n" + "=" * 80)
+ print("Generating response...")
+ print("=" * 80)
+
+ with torch.no_grad():
+ generated_ids = input_ids.clone()
+
+ for step in range(args.max_new_tokens):
+ # Forward pass
+ outputs = model(
+ input_ids=generated_ids,
+ image=pixel_values,
+ labels=None
+ )
+
+ # Get logits for next token
+ logits = outputs[0] # [batch, seq_len, vocab_size]
+ next_token_logits = logits[:, -1, :] / args.temperature
+
+ # Sample next token
+ probs = F.softmax(next_token_logits, dim=-1)
+ next_token = torch.multinomial(probs, num_samples=1)
+
+ # Append to sequence
+ generated_ids = torch.cat([generated_ids, next_token], dim=1)
+
+ # Check for EOS
+ if next_token.item() in [tokenizer.eos_token_id, 128009]: # llama3 eos
+ break
+
+ # Print progress
+ if (step + 1) % 10 == 0:
+ print(f" Generated {step + 1} tokens...")
+
+ # Decode output
+ output_ids = generated_ids[0, input_ids.shape[1]:] # Remove prompt
+ output_text = tokenizer.decode(output_ids, skip_special_tokens=True)
+
+ print("\n" + "=" * 80)
+ print("RESULTS")
+ print("=" * 80)
+ print(f"Prompt: {args.prompt}")
+ print(f"\nGenerated ({output_ids.shape[0]} tokens):")
+ print(output_text)
+ print("=" * 80)
+
+ print("\nโ Inference complete!")
+
+if __name__ == "__main__":
+ main()
diff --git a/aiak_megatron/megatron/core/datasets/Makefile b/aiak_megatron/megatron/core/datasets/Makefile
index e745f523..7e2344c7 100644
--- a/aiak_megatron/megatron/core/datasets/Makefile
+++ b/aiak_megatron/megatron/core/datasets/Makefile
@@ -7,7 +7,15 @@ LIBEXT = $(shell python3-config --extension-suffix)
OUT = $(LIBNAME)$(LIBEXT)
SRC = helpers.cpp
-default: $(OUT)
+# Check if any compiled library exists
+EXISTING_SO = $(wildcard helpers_cpp*.so)
+
+default:
+ifneq ($(EXISTING_SO),)
+ @echo "Using existing compiled library: $(EXISTING_SO)"
+else
+ @$(MAKE) $(OUT)
+endif
$(OUT): $(SRC)
$(CXX) $(CXXFLAGS) $(CPPFLAGS) $< -o $@
diff --git a/aiak_megatron/megatron/core/datasets/utils.py b/aiak_megatron/megatron/core/datasets/utils.py
index cb84c89a..abb57fea 100644
--- a/aiak_megatron/megatron/core/datasets/utils.py
+++ b/aiak_megatron/megatron/core/datasets/utils.py
@@ -22,8 +22,21 @@ def compile_helpers():
"""Compile C++ helper functions at runtime. Make sure this is invoked on a single process."""
import os
import subprocess
-
- command = ["make", "-C", os.path.abspath(os.path.dirname(__file__))]
+ import glob
+
+ # Check if helpers_cpp is already compiled
+ helpers_dir = os.path.abspath(os.path.dirname(__file__))
+ so_files = glob.glob(os.path.join(helpers_dir, "helpers_cpp*.so"))
+ if so_files:
+ # Try to import to verify it works
+ try:
+ import helpers_cpp
+ log_single_rank(logger, logging.INFO, f"Using pre-compiled helpers: {so_files[0]}")
+ return
+ except ImportError:
+ pass
+
+ command = ["make", "-C", helpers_dir]
if subprocess.run(command).returncode != 0:
import sys
diff --git a/aiak_megatron/megatron/core/dist_checkpointing/strategies/base.py b/aiak_megatron/megatron/core/dist_checkpointing/strategies/base.py
index 77eb325b..91b71106 100644
--- a/aiak_megatron/megatron/core/dist_checkpointing/strategies/base.py
+++ b/aiak_megatron/megatron/core/dist_checkpointing/strategies/base.py
@@ -23,6 +23,7 @@ class StrategyAction(Enum):
default_strategies: DefaultDict[str, dict[tuple, Any]] = defaultdict(dict)
+
async_calls = AsyncCallsQueue()
diff --git a/aiak_megatron/megatron/core/extensions/transformer_engine.py b/aiak_megatron/megatron/core/extensions/transformer_engine.py
index 82feb62c..c934ac87 100644
--- a/aiak_megatron/megatron/core/extensions/transformer_engine.py
+++ b/aiak_megatron/megatron/core/extensions/transformer_engine.py
@@ -1,302 +1,124 @@
-# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
-"""transformer engine utilities"""
-
-import dataclasses
-import io
-import os
-import pickle
-import warnings
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+"""
+TE-stub: Transformer Engine *disabled* version for Megatron-Core.
+
+This file exposes the same public symbols as the original
+`megatron.core.extensions.transformer_engine` module, but:
+
+ * Does NOT import `transformer_engine` at all.
+ * Provides simple PyTorch-based fallbacks for a few classes.
+ * Marks TE-specific / FP8 / advanced features as unavailable
+ and raises clear RuntimeError if they are actually used.
+
+The goal is that:
+ - `from megatron.core.extensions.transformer_engine import ...`
+ always succeeds, and
+ - as long as your config does NOT explicitly request TE/FP8
+ (e.g. you use bf16 and non-TE attention backend), training
+ can run without touching any real TE code.
+"""
+
+from contextlib import nullcontext
from typing import Any, Callable, Optional
import torch
-import transformer_engine as te
-from packaging.version import Version as PkgVersion
-from torch import Tensor
-from torch.nn.parameter import Parameter
-
-from megatron.core.dist_checkpointing.utils import replace_prefix_for_sharding
-from megatron.core.model_parallel_config import ModelParallelConfig
-from megatron.core.packed_seq_params import PackedSeqParams
-from megatron.core.parallel_state import (
- get_context_parallel_global_ranks,
- get_context_parallel_group,
- get_expert_data_parallel_rank,
- get_expert_model_parallel_rank,
- get_expert_model_parallel_world_size,
- get_expert_tensor_parallel_group,
- get_expert_tensor_parallel_rank,
- get_expert_tensor_parallel_world_size,
- get_hierarchical_context_parallel_groups,
- get_tensor_model_parallel_group,
- get_tensor_model_parallel_rank,
- get_tensor_model_parallel_world_size,
-)
-from megatron.core.tensor_parallel import get_cuda_rng_tracker, get_expert_parallel_rng_tracker_name
-from megatron.core.tensor_parallel.layers import (
- _initialize_affine_weight_cpu,
- set_tensor_model_parallel_attributes,
-)
-from megatron.core.tensor_parallel.random import get_data_parallel_rng_tracker_name
-from megatron.core.tensor_parallel.utils import divide
-from megatron.core.transformer.enums import AttnMaskType
-from megatron.core.enums import Fp8Recipe
-from megatron.core.transformer.transformer_config import TransformerConfig
-from megatron.core.transformer.utils import make_sharded_tensors_for_checkpoint
-from megatron.core.utils import get_te_version, is_te_min_version
-
-
-def _get_extra_te_kwargs(config: TransformerConfig):
- extra_transformer_engine_kwargs = {"params_dtype": config.params_dtype}
-
- if is_te_min_version("0.12.0"):
- if config.use_cpu_initialization:
- extra_transformer_engine_kwargs["device"] = 'cpu'
- elif config.init_model_with_meta_device:
- extra_transformer_engine_kwargs["device"] = "meta"
- else:
- extra_transformer_engine_kwargs["device"] = torch.cuda.current_device()
- return extra_transformer_engine_kwargs
+from torch import nn, Tensor
-def condition_init_method(config, init_method):
- """Condition TE init_method on config.perform_initialization."""
- return init_method if config.perform_initialization else (lambda w: None)
+# ---------------------------------------------------------------------------
+# Global TE capability flag & MoE fused permute stubs
+# ---------------------------------------------------------------------------
+HAVE_TE = False # other code may check this
-class TENorm:
- """
- A conditional wrapper to initialize an instance of Transformer-Engine's
- `LayerNorm` or `RMSNorm` based on input
- """
+# MoE fused permute helpers โ not available without TE
+fused_permute = None
+fused_unpermute = None
+fused_sort_chunks_by_index = None
- # TODO should we ditch normalization config and just use spec to choose LayerNorm vs RMSNorm?
- def __new__(cls, config: TransformerConfig, hidden_size: int, eps: float = 1e-5):
- if config.normalization == "LayerNorm":
- instance = te.pytorch.LayerNorm(
- hidden_size=hidden_size,
- eps=eps,
- sequence_parallel=config.sequence_parallel,
- zero_centered_gamma=config.layernorm_zero_centered_gamma,
- **_get_extra_te_kwargs(config),
- )
- elif config.normalization == "RMSNorm":
- assert hasattr(
- te.pytorch, "RMSNorm"
- ), "Transformer-Engine >= v0.11 required to use this feature"
- instance = te.pytorch.RMSNorm(
- hidden_size=hidden_size,
- eps=eps,
- sequence_parallel=config.sequence_parallel,
- zero_centered_gamma=config.layernorm_zero_centered_gamma,
- **_get_extra_te_kwargs(config),
- )
- else:
- raise Exception('Only LayerNorm and RMSNorm are curently supported')
+# FP8 padding helpers โ not available
+Fp8Padding = None
+Fp8Unpadding = None
- return instance
+# ---------------------------------------------------------------------------
+# Utility: simple checkpoint wrapper (no real TE checkpoint)
+# ---------------------------------------------------------------------------
-class TELinear(te.pytorch.Linear):
+def te_checkpoint(
+ forward_func: Callable,
+ distribute_saved_activations: bool,
+ get_rng_state_tracker: Any,
+ tp_group: Any,
+ hidden_states: Tensor,
+ attention_mask: Optional[Tensor],
+ attn_mask_type: Any,
+ context: Optional[Tensor],
+ context_mask: Optional[Tensor],
+ rotary_pos_emb: Optional[Tensor],
+ **kwargs,
+):
"""
- Wrapper for the Transformer-Engine's `Linear` layer.
-
- Note that if Megatron's parallel_state has not been initialized
- yet, the tp_group passed to TE will be None and must be set later
- via set_tensor_parallel_group().
-
- parallel_mode currently supports 3 different values:
- - "column": Split the weight matrix along output dimension (used in TEColumnParallelLinear)
- - "row": Split the weight matrix along input dimension (used in TERowParallelLinear)
- - "duplicated": No tensor parallelism and weight is duplicated across TP ranks
- - Note: For expert linear layers, we will disable communication logic here
- as TP communication is handled in token_dispatcher.
+ Minimal replacement for TE checkpointing.
+
+ We simply call `forward_func` directly WITHOUT activation checkpointing.
+ This keeps correctness but loses memory savings compared to real TE.
"""
+ return forward_func(
+ hidden_states,
+ attention_mask,
+ attn_mask_type,
+ context,
+ context_mask,
+ rotary_pos_emb,
+ **kwargs,
+ )
- def __init__(
- self,
- input_size: int,
- output_size: int,
- *,
- parallel_mode: Optional[str],
- config: ModelParallelConfig,
- init_method: Callable,
- bias: bool,
- skip_bias_add: bool,
- skip_weight_param_allocation: bool,
- tp_comm_buffer_name: Optional[str] = None,
- is_expert: bool = False,
- ):
- self.config = config
- # TE returns a zero length Tensor when bias=False and
- # return_bias=True, but we prefer None. So in that case we
- # tell TE to not return the bias, and return None
- # ourselves. This way our forward always returns two values
- # and we don't have to deal with the zero length Tensor.
- self.te_return_bias = skip_bias_add and bias
- self.is_first_microbatch = True
- self.disable_parameter_transpose_cache = self.config.disable_parameter_transpose_cache
- if skip_weight_param_allocation:
- raise ValueError(
- 'Transformer Engine linear layers do not support skip_weight_param_allocation'
- )
+# ---------------------------------------------------------------------------
+# TENorm โ fallback to plain LayerNorm
+# ---------------------------------------------------------------------------
- extra_kwargs = _get_extra_te_kwargs(config)
-
- if self.config.split_bw:
- extra_kwargs["delay_wgrad_compute"] = self.config.split_bw
-
- if (
- self.config.tp_comm_overlap
- and tp_comm_buffer_name
- and tp_comm_buffer_name not in ['qkv', 'proj', 'fc1', 'fc2']
- ):
- self.config.tp_comm_overlap = False
- warnings.warn(
- f"The user buffer name {tp_comm_buffer_name} is not supported in"
- "Transformer Engine. Disabling TP communication overlap "
- "for this layer."
- )
+class TENorm(nn.Module):
+ """
+ TE-style normalization wrapper.
- if is_te_min_version("0.8.0"):
- if self.config.tp_comm_overlap:
- if is_te_min_version("1.5.0"):
- # Use old overlap flags if they were supplied instead
- extra_kwargs["ub_overlap_ag"] = (
- self.config.tp_comm_overlap_ag
- if hasattr(self.config, "tp_comm_overlap_ag")
- else self.config.tp_comm_split_ag or self.config.tp_comm_atomic_ag
- )
- extra_kwargs["ub_overlap_rs"] = (
- self.config.tp_comm_overlap_rs
- if hasattr(self.config, "tp_comm_overlap_rs")
- else self.config.tp_comm_split_rs or self.config.tp_comm_atomic_rs
- )
- # Disable ub overlap for experts.
- if is_expert:
- extra_kwargs["ub_overlap_ag"] = False
- extra_kwargs["ub_overlap_rs"] = False
- else:
- extra_kwargs["ub_split_ag"] = self.config.tp_comm_split_ag
- extra_kwargs["ub_atomic_gemm_ag"] = self.config.tp_comm_atomic_ag
- extra_kwargs["ub_split_rs"] = self.config.tp_comm_split_rs
- extra_kwargs["ub_atomic_gemm_rs"] = self.config.tp_comm_atomic_rs
- # Disable ub overlap for experts.
- if is_expert:
- extra_kwargs["ub_split_ag"] = False
- extra_kwargs["ub_atomic_gemm_ag"] = False
- extra_kwargs["ub_split_rs"] = False
- extra_kwargs["ub_atomic_gemm_rs"] = False
- if is_te_min_version("1.0.0", check_equality=False):
- assert (
- tp_comm_buffer_name is not None
- ), "Buffer name should be set to configure communication overlap settings"
- extra_kwargs["ub_name"] = tp_comm_buffer_name
-
- self.expert_parallel = self.config.expert_model_parallel_size > 1
- if is_expert:
- rng_tracker_name = get_expert_parallel_rng_tracker_name()
- else:
- if parallel_mode == "duplicated":
- rng_tracker_name = get_data_parallel_rng_tracker_name()
- else:
- rng_tracker_name = None
- if is_te_min_version("1.7.0"):
- extra_kwargs["rng_tracker_name"] = rng_tracker_name
-
- te_parallel_mode = parallel_mode
- if parallel_mode == "duplicated":
- # Handle non-parallel case
- tp_group = None
- tp_size = 1
- explicit_expert_comm = False
- te_parallel_mode = None
+ Original TE version chooses between LayerNorm / RMSNorm etc.
+ Here we just always use plain LayerNorm.
+
+ Original signature:
+ TENorm(config: TransformerConfig, hidden_size: int, eps: float = 1e-5)
+ We accept arbitrary *args / **kwargs and try to pick hidden_size.
+ """
+
+ def __init__(self, *args, **kwargs):
+ super().__init__()
+
+ # Try to infer hidden_size and eps
+ if len(args) >= 2:
+ hidden_size = args[1]
else:
- # Disable communications in TE when using TP or EP by
- # making TE agnostic of model parallel.
- if is_expert:
- tp_group = get_expert_tensor_parallel_group(check_initialized=False)
- tp_size = get_expert_tensor_parallel_world_size()
- else:
- tp_group = get_tensor_model_parallel_group(check_initialized=False)
- tp_size = get_tensor_model_parallel_world_size()
- explicit_expert_comm = is_expert and (tp_size > 1 or self.expert_parallel)
+ hidden_size = kwargs.get("hidden_size") or kwargs.get("normalized_shape")
+ if hidden_size is None:
+ raise ValueError(
+ "TENorm stub could not infer hidden_size. "
+ "Expected signature: TENorm(config, hidden_size, eps=1e-5)"
+ )
- if explicit_expert_comm:
- if parallel_mode == "column":
- output_size = divide(output_size, tp_size)
- elif parallel_mode == "row":
- input_size = divide(input_size, tp_size)
- te_parallel_mode = None
- tp_size = 1
- tp_group = None
+ eps = kwargs.get("eps", 1e-5)
+ self.ln = nn.LayerNorm(hidden_size, eps=eps)
+ def forward(self, x: Tensor) -> Tensor:
+ return self.ln(x)
- super().__init__(
- in_features=input_size,
- out_features=output_size,
- sequence_parallel=self.config.sequence_parallel,
- fuse_wgrad_accumulation=self.config.gradient_accumulation_fusion,
- tp_group=tp_group,
- tp_size=tp_size,
- get_rng_state_tracker=(
- get_cuda_rng_tracker if get_cuda_rng_tracker().is_initialized() else None
- ),
- init_method=condition_init_method(config, init_method),
- bias=bias,
- return_bias=self.te_return_bias,
- parallel_mode=te_parallel_mode,
- **extra_kwargs,
- )
- for param in self.parameters():
- if is_expert:
- # Reduce the gradient on the expert_data_parallel group for expert linear layers
- setattr(param, 'allreduce', not self.expert_parallel)
- else:
- # Reduce the gradient on DP group
- setattr(param, 'allreduce', True)
- if parallel_mode == "duplicated":
- # Reduce the gradient further on the TP group since the weight is
- # duplicated across TP ranks
- setattr(param, 'sequence_parallel', self.config.sequence_parallel)
-
- def forward(self, x):
- """Forward."""
- _is_first_microbatch = (
- None if self.disable_parameter_transpose_cache else self.is_first_microbatch
- )
+# ---------------------------------------------------------------------------
+# TELinear and variants โ minimal PyTorch fallbacks
+# ---------------------------------------------------------------------------
- out = super().forward(x, is_first_microbatch=_is_first_microbatch)
- self.is_first_microbatch = False
-
- # TE only returns a tuple when return_bias is True, otherwise
- # it returns a single Tensor, we always want to return two
- # values regardless of the arguments.
- if self.te_return_bias:
- return out
- return out, None
-
- def sharded_state_dict(self, prefix='', sharded_offsets=(), metadata=None):
- """Replicate cross TP/DP."""
-
- # Provide the dist-ckpt support when TELinear is directly used
- # It can only happen with duplicated parallel mode
- assert (
- self.parallel_mode is None
- ), "TELinear sharded_state_dict can only be used with duplicated parallel mode"
- state_dict = self.state_dict(prefix='', keep_vars=True)
- return make_sharded_tensors_for_checkpoint(state_dict, prefix, None, sharded_offsets)
-
- def backward_dw(self):
- """Compute weight gradients during the backward pass if split_bw is enabled."""
- if self.config.split_bw:
- super().backward_dw()
-
-class TELayerNormColumnParallelLinear(te.pytorch.LayerNormLinear):
+class TELinear(nn.Module):
"""
- Wrapper for the Transformer-Engine's `LayerNormLinear` layer that combines
- layernorm and linear layers
+ Fallback for TE's Linear wrappers.
"""
def __init__(
@@ -304,174 +126,54 @@ def __init__(
input_size: int,
output_size: int,
*,
- config: TransformerConfig,
+ config: Any,
init_method: Callable,
- gather_output: bool,
bias: bool,
skip_bias_add: bool,
- is_expert: bool,
skip_weight_param_allocation: bool = False,
tp_comm_buffer_name: Optional[str] = None,
+ parallel_mode: Optional[str] = None,
+ is_expert: bool = False,
+ **kwargs,
):
- self.config = config
-
- if gather_output:
- raise ValueError('Transformer Engine linear layers do not support gather_output = True')
-
- if is_expert:
- raise ValueError('Transformer Engine linear layers do not yet support MoE')
-
+ super().__init__()
if skip_weight_param_allocation:
- raise ValueError(
- 'Transformer Engine linear layers do not support skip_weight_param_allocation'
+ raise RuntimeError(
+ "TELinear stub does not support skip_weight_param_allocation=True."
)
- # TE returns a zero length Tensor when bias=False and
- # return_bias=True, but we prefer None. So in that case we
- # tell TE to not return the bias, and return None
- # ourselves. This way our forward always returns two values
- # and we don't have to deal with the zero length Tensor.
- self.te_return_bias = skip_bias_add and bias
- self.is_first_microbatch = True
- self.disable_parameter_transpose_cache = self.config.disable_parameter_transpose_cache
- extra_kwargs = _get_extra_te_kwargs(config)
-
- if self.config.split_bw:
- extra_kwargs["delay_wgrad_compute"] = self.config.split_bw
-
- # Only Transformer-Engine version >= 0.11.0 supports `RMSNorm`
- if is_te_min_version("0.11.0"):
- extra_kwargs["normalization"] = self.config.normalization
- elif self.config.normalization != "LayerNorm":
- te_version = get_te_version()
- raise ValueError(
- f"Transformer Engine v{te_version} does not support {self.config.normalization}."
- )
-
- if is_te_min_version("0.8.0"):
- if self.config.tp_comm_overlap:
- extra_kwargs["ub_bulk_wgrad"] = self.config.tp_comm_bulk_wgrad
- extra_kwargs["ub_bulk_dgrad"] = self.config.tp_comm_bulk_dgrad
- if is_te_min_version("1.5.0", check_equality=False):
- # Use old overlap flags if they were supplied instead
- extra_kwargs["ub_overlap_ag"] = (
- self.config.tp_comm_overlap_ag
- if hasattr(self.config, "tp_comm_overlap_ag")
- else self.config.tp_comm_split_ag or self.config.tp_comm_atomic_ag
- )
- if is_te_min_version("1.6.0.dev0", check_equality=False):
- extra_kwargs["ub_overlap_rs_dgrad"] = (
- self.config.tp_comm_overlap_rs_dgrad
- if hasattr(self.config, "tp_comm_overlap_rs_dgrad")
- else False
- )
- if tp_comm_buffer_name == 'qkv' and self.config.tp_comm_overlap_disable_qkv:
- extra_kwargs["ub_overlap_ag"] = False
- extra_kwargs["ub_overlap_rs_dgrad"] = False
-
- if tp_comm_buffer_name == 'fc1' and self.config.tp_comm_overlap_disable_fc1:
- extra_kwargs["ub_overlap_ag"] = False
- extra_kwargs["ub_overlap_rs_dgrad"] = False
- else:
- extra_kwargs["ub_atomic_gemm_ag"] = self.config.tp_comm_atomic_ag
- extra_kwargs["ub_split_ag"] = self.config.tp_comm_split_ag
- if is_te_min_version("1.0.0", check_equality=False):
- assert (
- tp_comm_buffer_name is not None
- ), "Buffer name should be set to configure communication overlap settings"
- extra_kwargs["ub_name"] = tp_comm_buffer_name
-
-
- super().__init__(
- in_features=input_size,
- out_features=output_size,
- eps=self.config.layernorm_epsilon,
- sequence_parallel=self.config.sequence_parallel,
- fuse_wgrad_accumulation=self.config.gradient_accumulation_fusion,
- tp_group=get_tensor_model_parallel_group(check_initialized=False),
- tp_size=self.config.tensor_model_parallel_size,
- get_rng_state_tracker=(
- get_cuda_rng_tracker if get_cuda_rng_tracker().is_initialized() else None
- ),
- init_method=(
- condition_init_method(config, init_method)
- if not config.use_cpu_initialization
- else lambda w: None
- ),
- bias=bias,
- return_bias=self.te_return_bias,
- parallel_mode="column",
- return_layernorm_output=False,
- zero_centered_gamma=self.config.layernorm_zero_centered_gamma,
- **extra_kwargs,
- )
-
- world_size = get_tensor_model_parallel_world_size()
- rank = get_tensor_model_parallel_rank()
-
- if config.use_cpu_initialization:
- output_size_per_partition = divide(output_size, world_size)
- _ = _initialize_affine_weight_cpu(
- self.weight,
- output_size,
- input_size,
- output_size_per_partition,
- 0,
- init_method=condition_init_method(config, init_method),
- stride=1,
- return_master_weight=False,
- rank=rank,
- world_size=world_size,
- skip_set_tensor_parallel_attributes=True,
- )
- if bias:
- self.bias = Parameter(
- torch.empty(output_size_per_partition, dtype=config.params_dtype)
- )
- set_tensor_model_parallel_attributes(self.bias, True, 0, 1)
- with torch.no_grad():
- self.bias.zero_()
- setattr(self.bias, 'allreduce', True)
-
- def forward(self, x):
- """Forward."""
- _is_first_microbatch = (
- None if self.disable_parameter_transpose_cache else self.is_first_microbatch
- )
-
- out = super().forward(x, is_first_microbatch=_is_first_microbatch)
- self.is_first_microbatch = False
-
- # TE only returns a tuple when return_bias is True, otherwise
- # it returns a single Tensor, we always want to return two
- # values regardless of the arguments.
- if self.te_return_bias:
- return out
- return out, None
-
- def sharded_state_dict(self, prefix='', sharded_offsets=(), metadata=None):
- """Sharding along axis 0, bias sharded"""
- state_dict = self.state_dict(prefix='', keep_vars=True)
- return make_sharded_tensors_for_checkpoint(
- state_dict, prefix, {'weight': 0, 'bias': 0}, sharded_offsets
- )
-
- def __repr__(self):
- return (
- f"{type(self).__name__}(in_features={self.in_features}, "
- f"out_features={self.out_features}, bias={self.use_bias}, TP={self.tp_size})"
- )
-
- def backward_dw(self):
- """Compute weight gradients during the backward pass if split_bw is enabled."""
- if self.config.split_bw:
- super().backward_dw()
+ self.linear = nn.Linear(input_size, output_size, bias=bias)
+ self.skip_bias_add = skip_bias_add
+
+ if init_method is not None:
+ init_method(self.linear.weight)
+
+ def forward(self, x: Tensor):
+ # Auto-adjust in_features on first mismatch to match runtime input.
+ # This keeps training moving if TP sizing assumptions differ.
+ if x.size(-1) != self.linear.in_features:
+ try:
+ device = self.linear.weight.device
+ bias_flag = self.linear.bias is not None
+ out_features = self.linear.out_features
+ # Recreate linear with correct in_features
+ new_linear = nn.Linear(x.size(-1), out_features, bias=bias_flag).to(device)
+ self.linear = new_linear
+ print(f"[TE-Stub] Adjusted Linear in_features to {x.size(-1)} (was mismatch)", flush=True)
+ except Exception:
+ pass
+ # Ensure dtype/device match input to avoid BF16/FP32 matmul mismatch
+ if self.linear.weight.dtype != x.dtype or self.linear.weight.device != x.device:
+ self.linear = self.linear.to(device=x.device, dtype=x.dtype)
+ y = self.linear(x)
+ if self.skip_bias_add:
+ return y, self.linear.bias
+ return y, None
class TEColumnParallelLinear(TELinear):
"""
- Wrapper for the Transformer-Engine's `Linear` layer but specialized similar
- to megatron's `ColumnParallelLinear` layer.
+ Fallback version of TEColumnParallelLinear with basic TP-aware sizing.
"""
def __init__(
@@ -479,7 +181,7 @@ def __init__(
input_size: int,
output_size: int,
*,
- config: ModelParallelConfig,
+ config: Any,
init_method: Callable,
gather_output: bool,
bias: bool,
@@ -487,80 +189,44 @@ def __init__(
is_expert: bool,
skip_weight_param_allocation: bool = False,
tp_comm_buffer_name: Optional[str] = None,
+ input_is_parallel: bool = False,
+ **kwargs,
):
if gather_output:
- raise ValueError('Transformer Engine linear layers do not support gather_output = True')
+ raise RuntimeError(
+ "TEColumnParallelLinear stub does not support gather_output=True."
+ )
+ try:
+ from megatron.core import parallel_state as _ps
+ tp_world = _ps.get_tensor_model_parallel_world_size()
+ except Exception:
+ tp_world = 1
+ local_in = input_size
+ local_out = output_size
+ if tp_world and tp_world > 1:
+ if input_is_parallel:
+ assert input_size % tp_world == 0, "input_size must be divisible by TP size"
+ local_in = input_size // tp_world
+ assert output_size % tp_world == 0, "output_size must be divisible by TP size"
+ local_out = output_size // tp_world
super().__init__(
- input_size=input_size,
- output_size=output_size,
+ input_size=local_in,
+ output_size=local_out,
parallel_mode="column",
config=config,
- init_method=(
- condition_init_method(config, init_method)
- if not config.use_cpu_initialization
- else lambda w: None
- ),
+ init_method=init_method,
bias=bias,
skip_bias_add=skip_bias_add,
- is_expert=is_expert,
skip_weight_param_allocation=skip_weight_param_allocation,
tp_comm_buffer_name=tp_comm_buffer_name,
+ is_expert=is_expert,
)
- if config.use_cpu_initialization:
- if is_expert:
- world_size = get_expert_tensor_parallel_world_size()
- rank = get_expert_tensor_parallel_rank()
- else:
- world_size = get_tensor_model_parallel_world_size()
- rank = get_tensor_model_parallel_rank()
- output_size_per_partition = divide(output_size, world_size)
- _ = _initialize_affine_weight_cpu(
- self.weight,
- output_size,
- input_size,
- output_size_per_partition,
- 0,
- init_method=condition_init_method(config, init_method),
- stride=1,
- return_master_weight=False,
- rank=rank,
- world_size=world_size,
- skip_set_tensor_parallel_attributes=True,
- )
- if bias:
- self.bias = Parameter(
- torch.empty(output_size_per_partition, dtype=config.params_dtype)
- )
- set_tensor_model_parallel_attributes(self.bias, True, 0, 1)
- with torch.no_grad():
- self.bias.zero_()
- setattr(self.bias, 'allreduce', True)
-
- def sharded_state_dict(self, prefix='', sharded_offsets=(), metadata=None):
- """Sharding along axis 0, bias sharded"""
- state_dict = self.state_dict(prefix='', keep_vars=True)
- return make_sharded_tensors_for_checkpoint(
- state_dict, prefix, {'weight': 0, 'bias': 0}, sharded_offsets
- )
-
- def __repr__(self):
- return (
- f"{type(self).__name__}(in_features={self.in_features}, "
- f"out_features={self.out_features}, bias={self.use_bias}, TP={self.tp_size})"
- )
-
- def backward_dw(self):
- """Compute weight gradients during the backward pass if split_bw is enabled."""
- if self.config.split_bw:
- super().backward_dw()
-
class TERowParallelLinear(TELinear):
"""
- Wrapper for the Transformer-Engine's `Linear` layer but specialized similar
- to megatron's `RowParallelLinear` layer.
+ Fallback version of TERowParallelLinear with basic TP-aware sizing.
"""
def __init__(
@@ -568,101 +234,179 @@ def __init__(
input_size: int,
output_size: int,
*,
- config: ModelParallelConfig,
+ config: Any,
init_method: Callable,
bias: bool,
input_is_parallel: bool,
skip_bias_add: bool,
is_expert: bool,
tp_comm_buffer_name: Optional[str] = None,
+ **kwargs,
):
if not input_is_parallel:
- raise ValueError(
- "Transformer Engine linear layers do not support input_is_parallel = False"
+ raise RuntimeError(
+ "TERowParallelLinear stub expects input_is_parallel=True (same as TE)."
)
-
+ try:
+ from megatron.core import parallel_state as _ps
+ tp_world = _ps.get_tensor_model_parallel_world_size()
+ except Exception:
+ tp_world = 1
+ local_in = input_size
+ local_out = output_size
+ if tp_world and tp_world > 1:
+ assert input_size % tp_world == 0, "input_size must be divisible by TP size"
+ local_in = input_size // tp_world
+ # Row-parallel returns full-sized outputs
super().__init__(
- input_size=input_size,
- output_size=output_size,
+ input_size=local_in,
+ output_size=local_out,
parallel_mode="row",
config=config,
- init_method=(
- condition_init_method(config, init_method)
- if not config.use_cpu_initialization
- else lambda w: None
- ),
+ init_method=init_method,
bias=bias,
skip_bias_add=skip_bias_add,
- skip_weight_param_allocation=False, # We don't currently use this for row parallel layers
- is_expert=is_expert,
+ skip_weight_param_allocation=False,
tp_comm_buffer_name=tp_comm_buffer_name,
+ is_expert=is_expert,
)
- if config.use_cpu_initialization:
- if is_expert:
- world_size = get_expert_tensor_parallel_world_size()
- rank = get_expert_tensor_parallel_rank()
- else:
- world_size = get_tensor_model_parallel_world_size()
- rank = get_tensor_model_parallel_rank()
- input_size_per_partition = divide(input_size, world_size)
- self.master_weight = _initialize_affine_weight_cpu(
- self.weight,
- output_size,
- input_size,
- input_size_per_partition,
- 1,
- init_method=condition_init_method(config, init_method),
- stride=1,
- return_master_weight=False,
- params_dtype=config.params_dtype,
- rank=rank,
- world_size=world_size,
- skip_set_tensor_parallel_attributes=True,
+
+
+class TELayerNormColumnParallelLinear(nn.Module):
+ """
+ Fallback for TELayerNormColumnParallelLinear.
+
+ Original TE combines LayerNorm + ColumnParallelLinear.
+ Here we implement: y = Linear(LayerNorm(x))
+
+ Forward returns (y, bias_or_none) to match TE behaviour.
+ """
+
+ def __init__(
+ self,
+ input_size: int,
+ output_size: int,
+ *,
+ config: Any,
+ init_method: Callable,
+ gather_output: bool,
+ bias: bool,
+ skip_bias_add: bool,
+ is_expert: bool,
+ skip_weight_param_allocation: bool = False,
+ tp_comm_buffer_name: Optional[str] = None,
+ **kwargs,
+ ):
+ super().__init__()
+ if gather_output:
+ raise RuntimeError(
+ "TELayerNormColumnParallelLinear stub does not support gather_output=True."
+ )
+ if is_expert:
+ # You could still allow this, but keep consistent with original constraint
+ raise RuntimeError(
+ "TELayerNormColumnParallelLinear stub does not support is_expert=True."
)
- if bias:
- self.bias = Parameter(torch.empty(output_size, dtype=config.params_dtype))
- # Always initialize bias to zero.
- with torch.no_grad():
- self.bias.zero_()
- setattr(self.bias, 'allreduce', True)
- setattr(self.bias, 'sequence_parallel', config.sequence_parallel)
-
- def sharded_state_dict(self, prefix='', sharded_offsets=(), metadata=None):
- """Sharding along axis 1, bias not sharded"""
- state_dict = self.state_dict(prefix='', keep_vars=True)
- return make_sharded_tensors_for_checkpoint(
- state_dict, prefix, {'weight': 1}, sharded_offsets
- )
- def __repr__(self):
- return (
- f"{type(self).__name__}(in_features={self.in_features}, "
- f"out_features={self.out_features}, bias={self.use_bias}, TP={self.tp_size})"
- )
+ eps = getattr(config, "layernorm_epsilon", 1e-5)
+ self.ln = nn.LayerNorm(input_size, eps=eps)
+ self.linear = nn.Linear(input_size, output_size, bias=bias)
+ self.skip_bias_add = skip_bias_add
- def backward_dw(self):
- """Compute weight gradients during the backward pass if split_bw is enabled."""
- if self.config.split_bw:
- super().backward_dw()
+ if init_method is not None:
+ init_method(self.linear.weight)
+
+ def forward(self, x: Tensor):
+ y = self.linear(self.ln(x))
+ if self.skip_bias_add:
+ return y, self.linear.bias
+ else:
+ return y, None
-class TEDotProductAttention(te.pytorch.DotProductAttention):
+# ---------------------------------------------------------------------------
+# GroupedLinear stubs โ names exist, but not implemented
+# ---------------------------------------------------------------------------
+
+class TEGroupedLinear(nn.Module):
+ """
+ Stub for TEGroupedLinear.
+
+ If your config ever tries to instantiate this, you are relying on
+ Transformer Engine's grouped FP8 GEMMs, which are not available
+ in this stubbed build.
"""
- Wrapper for the Transformer-Engine's `DotProductAttention` layer that also
- has "flash attention" enabled.
- Note that if Megatron's parallel_state has not been initialized yet, the
- tp_group and cp_group passed to TE will be None and must be set later
- via set_tensor_parallel_group() and set_context_parallel_group().
+ def __init__(self, *args, **kwargs):
+ super().__init__()
+ raise RuntimeError("TEGroupedLinear is not available (Transformer Engine disabled).")
+
+ def forward(self, *args, **kwargs):
+ raise RuntimeError("TEGroupedLinear is not available (Transformer Engine disabled).")
+
+
+class TEColumnParallelGroupedLinear(TEGroupedLinear):
+ pass
+
+
+class TERowParallelGroupedLinear(TEGroupedLinear):
+ pass
+
+
+# ---------------------------------------------------------------------------
+# Attention wrapper โ disabled (we use non-TE backend instead)
+# ---------------------------------------------------------------------------
+# class TEDotProductAttention(nn.Module):
+# """
+# Very simple scaled dot-product attention fallback.
+# This is only meant to keep imports and basic training running,
+# not for maximum performance.
+# """
+
+# def __init__(self, *args, **kwargs):
+# super().__init__()
+
+# def forward(
+# self,
+# query: torch.Tensor,
+# key: torch.Tensor,
+# value: torch.Tensor,
+# *args,
+# **kwargs,
+# ) -> torch.Tensor:
+# print(type(query), query)
+# d_k = query.size(-1)
+# scores = torch.matmul(query, key.transpose(-2, -1)) / (d_k ** 0.5)
+# attn = torch.softmax(scores, dim=-1)
+# return torch.matmul(attn, value)
+
+import math
+from typing import Optional, Any
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+# Keep the same name so it can be used as a drop-in
+class TEDotProductAttention(nn.Module):
+ """
+ Simple, readable scaled dot-product attention replacement for Transformer-Engine's
+ DotProductAttention. Meant as a fallback for correctness / debugging, not performance.
+ Forward signature kept compatible with the TE wrapper used previously.
+
+ Notes:
+ - Supports qkv_format 'sbhd' (seq, batch, head, dim) and 'bshd' (batch, seq, head, dim).
+ - Also accepts 3-D (B, S, D) inputs and will split heads automatically.
+ - attention_mask can be boolean (True=keep) or additive (float with -inf on masked positions)
+ - attention_bias (if provided) is added to attention scores before softmax.
"""
- cp_stream: torch.cuda.Stream = None
+ cp_stream: torch.cuda.Stream = None # keep attribute so callers referencing it won't fail
def __init__(
self,
- config: TransformerConfig,
+ config: Any,
layer_number: int,
- attn_mask_type: AttnMaskType,
+ attn_mask_type: Any,
attention_type: str,
attention_dropout: Optional[float] = None,
softmax_scale: Optional[float] = None,
@@ -670,783 +414,579 @@ def __init__(
v_channels: Optional[int] = None,
cp_comm_type: str = "p2p",
):
+ super().__init__()
+ # Store config values we may need
self.config = config
- self.te_forward_mask_type = False
- self.qkv_format: str = 'sbhd'
-
- if self.config.apply_query_key_layer_scaling != bool(
- int(os.getenv('NVTE_APPLY_QK_LAYER_SCALING', '0'))
- ):
- raise ValueError(
- f"apply_query_key_layer_scaling is {self.config.apply_query_key_layer_scaling} "
- f"but environment variable NVTE_APPLY_QK_LAYER_SCALING is "
- f"{os.getenv('NVTE_APPLY_QK_LAYER_SCALING')}. Transformer Engine does not support "
- f"setting query key layer scaling via argument, so these two must match."
- )
-
- extra_kwargs: dict[str, Any] = {}
- if is_te_min_version("0.11.0"):
- extra_kwargs["num_gqa_groups"] = self.config.num_query_groups
- elif self.config.num_query_groups != self.config.num_attention_heads:
- raise ValueError(
- f"Transformer Engine v{get_te_version()} does not support Grouped Query Attention, "
- f"use a newer version of Transformer Engine. "
- f"(num_query_groups ({self.config.num_query_groups}) != "
- f"num_attention_heads ({self.config.num_attention_heads}))"
- )
-
- if is_te_min_version("0.10.0"):
- extra_kwargs["attention_type"] = attention_type
- # older version don't need attention_type
-
- if is_te_min_version("0.12.0", check_equality=False):
- self.te_forward_mask_type = True
-
- # This check is important as CP config can be disabled while having a valid CP group
- # Example - Disabling CP for encoder while a valid CP group exists for decoder
- if self.config.context_parallel_size > 1:
- assert is_te_min_version("1.0.0"), "Only Transformer-Engine version >= 1.0.0 supports context parallelism!"
- if getattr(TEDotProductAttention, "cp_stream") is None:
- TEDotProductAttention.cp_stream = torch.cuda.Stream()
-
- extra_kwargs["cp_group"] = get_context_parallel_group(check_initialized=False)
- extra_kwargs["cp_global_ranks"] = get_context_parallel_global_ranks(check_initialized=False)
- extra_kwargs["cp_stream"] = TEDotProductAttention.cp_stream
-
- if is_te_min_version("1.10.0"):
- if cp_comm_type is None:
- extra_kwargs["cp_comm_type"] = "p2p"
-
- elif cp_comm_type == "a2a+p2p":
- assert is_te_min_version("1.12.0"), (
- f"Transformer-Engine v{get_te_version()} must be >= 1.12.0 to support"
- "hierarchical cp commucation."
- )
- extra_kwargs["cp_comm_type"] = "a2a+p2p"
- extra_kwargs["cp_group"] = get_hierarchical_context_parallel_groups(check_initialized=False)
- else:
- extra_kwargs["cp_comm_type"] = cp_comm_type
-
- if self.config.deterministic_mode:
- if int(os.getenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1")) != 0:
- raise RuntimeError(
- "deterministic_mode is on and we are using DotProductAttention from "
- "Transformer Engine, but NVTE_ALLOW_NONDETERMINISTIC_ALGO is not 0. "
- f"Currently set to: {os.getenv('NVTE_ALLOW_NONDETERMINISTIC_ALGO', 'not set')}."
+ # default qkv format preserved from your original class
+ self.qkv_format: str = "sbhd"
+ # num heads convenience
+ self.num_heads = getattr(config, "num_attention_heads", 1)
+ # total kv channels (if present on config). We'll use this only when needed to infer head dim
+ self.kv_channels = getattr(config, "kv_channels", None)
+ # optional dropout on attention weights (kept None unless provided)
+ self.attn_dropout = nn.Dropout(attention_dropout) if attention_dropout and attention_dropout > 0.0 else None
+ # allow externally-configured softmax scale (float) OR compute as 1/sqrt(head_dim) later
+ self.softmax_scale = softmax_scale
+ # keep set of packed seq param names to remain compatible with callers (ignored here)
+ self.kept_packed_seq_params = set()
+
+ def _ensure_4d_and_split_heads(self, x: Tensor) -> (Tensor, bool):
+ """
+ Ensure x is in (B, S, H, D_head) layout and return (x4d, was_3d)
+ Accepts:
+ - x shaped (S, B, H, D) if qkv_format == 'sbhd' (seq, batch, head, dim)
+ - x shaped (B, S, H, D) if qkv_format == 'bshd'
+ - x shaped (B, S, D) -> will split last dim into (H, D_head)
+ Returns x in (B, S, H, D_head) format and a flag indicating whether original was 3d.
+ """
+ was_3d = False
+ if x is None:
+ raise ValueError("TEDotProductAttention: input tensor is None")
+ if x.dim() == 4:
+ if self.qkv_format == "sbhd":
+ # (S, B, H, D) -> convert to (B, S, H, D)
+ x4 = x.transpose(0, 1).contiguous()
+ elif self.qkv_format == "bshd":
+ x4 = x.contiguous()
+ else:
+ # fallback: assume bshd if unknown
+ x4 = x.contiguous()
+ elif x.dim() == 3:
+ # (B, S, D) -> split D into (H, D_head)
+ was_3d = True
+ B, S, D = x.shape
+ if D % self.num_heads != 0:
+ raise ValueError(
+ f"TEDotProductAttention: embedding dim ({D}) is not divisible by num_heads ({self.num_heads})."
)
-
- if config.window_size is not None:
- # Check version
- assert is_te_min_version("1.2.0"), (
- f"Transformer-Engine v{get_te_version()} must be >= 1.2.0 to support sliding window attention."
- )
- extra_kwargs['window_size'] = config.window_size
-
- if is_te_min_version("1.10.0"):
- # TE 1.10.0 introduces the ability to set the different k and v channels
- kv_channels = (
- (k_channels, v_channels)
- if k_channels is not None and v_channels is not None
- else self.config.kv_channels
- )
- extra_kwargs['softmax_scale'] = softmax_scale
+ D_head = D // self.num_heads
+ x4 = x.view(B, S, self.num_heads, D_head).contiguous()
else:
- kv_channels = self.config.kv_channels
-
- self.kept_packed_seq_params = set(field.name for field in dataclasses.fields(PackedSeqParams))
- if get_te_version() < PkgVersion("1.3.0"):
- # TE 1.3.0 introduces precomputing max_seqlen to remove unnecessary kernels and D2H
- # copies (#555)
- # These two arguments did not exist prior to 1.3.0
- self.kept_packed_seq_params.discard("max_seqlen_q")
- self.kept_packed_seq_params.discard("max_seqlen_kv")
-
- if get_te_version() < PkgVersion("1.10.0"):
- # TE 1.8.0 introduces cu_seqlens_padded which is the cu_seqlens with paddings counted
- # in each individual sequence in THD format dataset
- # These two arguments did not exist prior to 1.8.0. Full support added in 1.10.0 (#1012)
- self.kept_packed_seq_params.discard("cu_seqlens_q_padded")
- self.kept_packed_seq_params.discard("cu_seqlens_kv_padded")
+ raise ValueError(f"Tensors must be 3D (B,S,D) or 4D. Got shape {tuple(x.shape)}")
+ return x4, was_3d
- super().__init__(
- num_attention_heads=self.config.num_attention_heads,
- kv_channels=kv_channels,
- attention_dropout=self.config.attention_dropout if attention_dropout is None else attention_dropout,
- attn_mask_type=attn_mask_type.name,
- sequence_parallel=self.config.sequence_parallel,
- tp_size=self.config.tensor_model_parallel_size,
- get_rng_state_tracker=(
- get_cuda_rng_tracker if get_cuda_rng_tracker().is_initialized() else None
- ),
- tp_group=get_tensor_model_parallel_group(check_initialized=False),
- layer_number=layer_number,
- **extra_kwargs,
- )
+ def _reconstruct_output(self, out4: Tensor, was_3d: bool) -> Tensor:
+ """
+ out4 is (B, S, H, D_head).
+ If was_3d True -> return (B, S, D) by merging head dim
+ If qkv_format previously was 'sbhd' and user originally passed sbhd (we tracked earlier),
+ we attempt to return in original layout. However to keep behavior simple we return:
+ - (B,S,D) if original inputs were 3D
+ - (B,S,H,D) otherwise
+ """
+ if was_3d:
+ B, S, H, Dh = out4.shape
+ return out4.view(B, S, H * Dh)
+ else:
+ return out4
def forward(
self,
query: Tensor,
key: Tensor,
value: Tensor,
- attention_mask: Tensor,
- attn_mask_type: AttnMaskType,
- attention_bias: Tensor = None,
- packed_seq_params: PackedSeqParams = None,
- ):
- """Forward."""
- packed_seq_kwargs = (
- {key: getattr(packed_seq_params, key) for key in self.kept_packed_seq_params}
- if packed_seq_params is not None
- else {}
- )
- # overwrite self.qkv_format depending on self.config.apply_rope_fusion, which can be set
- # after init
- if self.config.apply_rope_fusion and is_te_min_version("0.13.0", check_equality=False):
- self.qkv_format = 'bshd'
-
- qkv_format = packed_seq_kwargs.get('qkv_format', self.qkv_format)
-
- # WAR for peak memory usage.
- # See https://gitlab-master.nvidia.com/ADLR/megatron-lm/-/merge_requests/2388
- if self.config.apply_rope_fusion and qkv_format == 'bshd':
- query, key, value = [x.transpose(0, 1).contiguous() for x in (query, key, value)]
- # In PyTorch, the following two tensors are in fact the same:
- # Tensor with shape (1, S, H, D) and stride (S*H*D, H*D, D, 1)
- # Tensor with shape (1, S, H, D) and stride (H*D, H*D, D, 1)
- # Stride for a dimension that is 1 has no meaning, so tensors created two different ways
- # can have same shape but different strides.
- # We unify them to the first one to pass the stride check in TE
- if value.shape == key.shape and value.shape[0] == 1 and value.stride() != key.stride():
- value = value.as_strided(value.shape, key.stride())
-
- attention_bias_kwargs = {}
- if attention_bias is not None:
- assert is_te_min_version("1.2.0"), (
- f"Transformer-Engine v{get_te_version()} must be >= 1.2.0 to support `attention_bias`."
- )
- attention_bias_kwargs = dict(
- core_attention_bias_type='post_scale_bias', core_attention_bias=attention_bias
- )
- else:
- if self.config.position_embedding_type == "alibi":
- attention_bias_kwargs = dict(core_attention_bias_type='alibi', core_attention_bias=None)
-
- if self.te_forward_mask_type:
- if qkv_format == 'thd' and is_te_min_version("1.7.0"):
- # thd format uses flash attention with cuDNN kernel which requires is_padding=True,
- # so the only acceptable mask types are `padding_causal` and `padding`. These do not
- # necessarily indicate there are padded tokens in the sequence.
- if attn_mask_type == AttnMaskType.causal:
- attn_mask_type = AttnMaskType.padding_causal
- elif attn_mask_type == AttnMaskType.no_mask:
- attn_mask_type = AttnMaskType.padding
- core_attn_out = super().forward(
- query,
- key,
- value,
- attention_mask,
- attn_mask_type=attn_mask_type.name,
- **attention_bias_kwargs,
- **packed_seq_kwargs,
- )
- else:
- core_attn_out = super().forward(
- query,
- key,
- value,
- attention_mask,
- **attention_bias_kwargs,
- **packed_seq_kwargs,
- )
-
- if self.config.apply_rope_fusion and qkv_format == 'bshd':
- return core_attn_out.transpose(0, 1)
- else:
- return core_attn_out
-
-
-if is_te_min_version("1.9.0.dev0"):
-
- class TEGroupedLinear(te.pytorch.GroupedLinear):
+ attention_mask: Optional[Tensor],
+ attn_mask_type: Any,
+ attention_bias: Optional[Tensor] = None,
+ packed_seq_params: Any = None,
+ ) -> Tensor:
"""
- Wrapper for the Transformer-Engine's `GroupedLinear` layer.
+ Simple scaled dot-product attention:
+ out[b, s_q, h, d] = softmax( (q*k^T)/sqrt(d) + attention_bias + mask ) @ v
- Note that if Megatron's parallel_state has not been initialized
- yet, the tp_group passed to TE will be None and must be set later
- via set_tensor_parallel_group().
+ Returns output in (B, S_q, H, D_head) unless the inputs were 3D (B,S,D) in which case returns (B,S,D).
"""
- def __init__(
- self,
- num_gemms: int,
- input_size: int,
- output_size: int,
- *,
- parallel_mode: Optional[str],
- config: ModelParallelConfig,
- init_method: Callable,
- bias: bool,
- skip_bias_add: bool,
- is_expert: bool = False,
- tp_comm_buffer_name: Optional[str] = None,
- ):
- self.config = config
-
- # TE returns a zero length Tensor when bias=False and
- # return_bias=True, but we prefer None. So in that case we
- # tell TE to not return the bias, and return None
- # ourselves. This way our forward always returns two values
- # and we don't have to deal with the zero length Tensor.
- self.te_return_bias = skip_bias_add and bias
- self.is_first_microbatch = True
- self.disable_parameter_transpose_cache = self.config.disable_parameter_transpose_cache
-
- extra_kwargs = _get_extra_te_kwargs(config)
-
- if self.config.split_bw:
- extra_kwargs["delay_wgrad_compute"] = self.config.split_bw
-
- extra_kwargs["ub_name"] = tp_comm_buffer_name
-
- self.expert_parallel = self.config.expert_model_parallel_size > 1
- if is_expert:
- extra_kwargs["rng_tracker_name"] = get_expert_parallel_rng_tracker_name()
-
- # The comms between TP and EP group is explicitly handled by MoE token dispatcher.
- # So we disable comms by making TE agnostic of model parallel.
- if is_expert:
- tp_group = get_expert_tensor_parallel_group(check_initialized=False)
- tp_size = get_expert_tensor_parallel_world_size()
+ # Allow callers to pass packed_seq_params (kept for API compatibility) but we ignore it here.
+ _ = packed_seq_params # noqa: F841
+
+ # Convert inputs to (B, S, H, D_head)
+ q4, q_was_3d = self._ensure_4d_and_split_heads(query)
+ k4, k_was_3d = self._ensure_4d_and_split_heads(key)
+ v4, v_was_3d = self._ensure_4d_and_split_heads(value)
+
+ # sanity check shapes
+ if not (q_was_3d == k_was_3d == v_was_3d):
+ # mixing 3D / 4D forms is confusing; require consistent input forms
+ raise ValueError("TEDotProductAttention: query/key/value must be all 3D or all 4D (consistent formats).")
+
+ Bq, Sq, Hq, Dh = q4.shape
+ Bk, Sk, Hk, Dk = k4.shape
+ Bv, Sv, Hv, Dv = v4.shape
+
+ if not (Bq == Bk == Bv):
+ raise ValueError("Batch sizes of query/key/value must match")
+ if not (Hq == Hk == Hv):
+ raise ValueError("Number of heads of query/key/value must match")
+ if not (Dk == Dv):
+ raise ValueError("Key/value head dims must match")
+ if Dh != Dk:
+ # head dims should match for q and k
+ # if they don't, we can still compute by projecting; but here we require equality
+ raise ValueError(f"Head dimension mismatch: q head dim {Dh} vs k head dim {Dk}")
+
+ B = Bq
+ # compute scale
+ scale = self.softmax_scale if self.softmax_scale is not None else 1.0 / math.sqrt(Dh)
+
+ # compute raw scores: we want shape (B, H, Sq, Sk)
+ # reshape to (B, H, Sq, Dh) etc and use matmul after transposing
+ # q4: (B, Sq, H, Dh) -> (B, H, Sq, Dh)
+ qBH = q4.transpose(1, 2) # (B, H, Sq, Dh)
+ kBH = k4.transpose(1, 2) # (B, H, Sk, Dh)
+ vBH = v4.transpose(1, 2) # (B, H, Sv, Dh)
+
+ # scores: (B, H, Sq, Sk)
+ scores = torch.matmul(qBH, kBH.transpose(-2, -1)) * scale
+
+ # apply attention_bias if given (broadcastable)
+ if attention_bias is not None:
+ # attention_bias commonly has shape (B, 1, 1, Sk) or (1, 1, 1, Sk) or (Sq, Sk)
+ scores = scores + attention_bias
+
+ # apply attention_mask
+ if attention_mask is not None:
+ # If mask is boolean where True means keep, convert to additive mask
+ # Accept both boolean masks and additive masks (float)
+ if attention_mask.dtype == torch.bool:
+ # mask shape might be (B, Sk) or (B, 1, 1, Sk) etc; we want True where keep
+ # we fill disallowed positions with -inf
+ scores = scores.masked_fill(~attention_mask, float("-inf"))
else:
- tp_group = get_tensor_model_parallel_group(check_initialized=False)
- tp_size = get_tensor_model_parallel_world_size()
- self.explicit_expert_comm = is_expert and (tp_size > 1 or self.expert_parallel)
-
- if self.explicit_expert_comm:
- if parallel_mode == "column":
- output_size = divide(output_size, tp_size)
- elif parallel_mode == "row":
- input_size = divide(input_size, tp_size)
- parallel_mode = None
- tp_size = 1
- tp_group = None
-
-
- super().__init__(
- num_gemms=num_gemms,
- in_features=input_size,
- out_features=output_size,
- sequence_parallel=self.config.sequence_parallel,
- fuse_wgrad_accumulation=self.config.gradient_accumulation_fusion,
- tp_group=tp_group,
- tp_size=tp_size,
- get_rng_state_tracker=(
- get_cuda_rng_tracker if get_cuda_rng_tracker().is_initialized() else None
- ),
- init_method=condition_init_method(config, init_method),
- bias=bias,
- return_bias=self.te_return_bias,
- parallel_mode=parallel_mode,
- **extra_kwargs,
- )
+ # assume additive mask already shaped for broadcasting
+ scores = scores + attention_mask
- for param in self.parameters():
- setattr(param, 'allreduce', not (is_expert and self.expert_parallel))
-
- def merge_extra_states(
- self,
- state_dict,
- prefix,
- local_metadata,
- strict,
- missing_keys,
- unexpected_keys,
- error_msgs,
- ):
- """
- Merge multiple "_extra_state" into one.
- """
- self.init_fp8_metadata(num_gemms=self.num_gemms)
- # When resume training, loading ckpt is out of fp8_autocast context.
- # So we need to manually detect from the state_dict.
- fp8_checkpoint = any("_extra_state" in str(key) for key in state_dict.keys())
-
- if not fp8_checkpoint:
- return
-
- try:
- state_list = [
- state_dict.pop(f"{prefix}_extra_state{i}") for i in range(1, self.num_gemms)
- ]
- except KeyError:
- # "_extra_state{i}" only exists for dist-ckpt. Return for torch native ckpt.
- return
-
- # Early return conditions:
- # 1. Empty state_dict
- # 2. Empty state_list
- # 3. _extra_state is None
- # 4. _extra_state does not contain any information
- if (
- not state_dict
- or not state_list
- or state_dict.get(f"{prefix}_extra_state") is None
- or self._decode_extra_state(state_dict[f"{prefix}_extra_state"]) is None
- ):
- return
-
- state_list = [state_dict.pop(f"{prefix}_extra_state")] + state_list
- state_list = [self._decode_extra_state(state) for state in state_list]
- extra_fp8_variables = state_list[0]['extra_fp8_variables']
- extra_fp8_variables['num_gemms'] = self.num_gemms
- extra_state = {"extra_fp8_variables": extra_fp8_variables}
- # TE 2.0 adds recipe in extra_state
- if is_te_min_version("2.0.0"):
- self.fp8_meta["recipe"] = state_list[0]['recipe']
- extra_state['recipe'] = self.fp8_meta["recipe"]
- # Only delayed scaling has global fp8 meta tensors. We're not using
- # self.fp8_meta["recipe"].delayed() because it's available in TE 2.0 and later.
- if isinstance(self.fp8_meta["recipe"], te.common.recipe.DelayedScaling):
- extra_state.update(
- {
- "scale_fwd": torch.cat(
- [state['scale_fwd'].view(-1, 1) for state in state_list], dim=1
- ).view(-1),
- "amax_history_fwd": torch.cat(
- [state['amax_history_fwd'].view(-1, 1) for state in state_list],
- dim=1,
- ).view(self.fp8_meta["recipe"].amax_history_len, -1),
- "scale_bwd": torch.cat(
- [state['scale_bwd'].view(-1, 1) for state in state_list], dim=1
- ).view(-1),
- "amax_history_bwd": torch.cat(
- [state['amax_history_bwd'].view(-1, 1) for state in state_list],
- dim=1,
- ).view(self.fp8_meta["recipe"].amax_history_len, -1),
- }
- )
- # TE 2.0 removes scale_inv_fwd and scale_inv_bwd
- if not is_te_min_version("2.0.0"):
- extra_state.update(
- {
- "scale_inv_fwd": torch.cat(
- [state['scale_inv_fwd'].view(-1, 1) for state in state_list],
- dim=1,
- ).view(-1),
- "scale_inv_bwd": torch.cat(
- [state['scale_inv_bwd'].view(-1, 1) for state in state_list],
- dim=1,
- ).view(-1),
- }
- )
- state_dict[f"{prefix}_extra_state"] = self._encode_extra_state(extra_state)
-
- self._register_load_state_dict_pre_hook(merge_extra_states, with_module=True)
-
- def forward(self, x, m_splits):
- """Forward."""
- _is_first_microbatch = (
- None if self.disable_parameter_transpose_cache else self.is_first_microbatch
- )
- out = super().forward(x, m_splits, is_first_microbatch=_is_first_microbatch)
- self.is_first_microbatch = False
-
- # TE only returns a tuple when return_bias is True, otherwise
- # it returns a single Tensor, we always want to return two
- # values regardless of the arguments.
- if self.te_return_bias:
- return out
- return out, None
-
- def _encode_extra_state(self, state):
- # TE 2.0 changed the format of extra_state to be a byte tensor
- if is_te_min_version("2.0.0"):
- torch.cuda.synchronize()
- state_serialized = bytearray(pickle.dumps(state))
- state_serialized = torch.frombuffer(state_serialized, dtype=torch.uint8)
- else:
- state_serialized = io.BytesIO()
- torch.save(state, state_serialized)
- return state_serialized
-
- def _decode_extra_state(self, state):
- if isinstance(state, torch.Tensor):
- return pickle.loads(state.detach().cpu().numpy().tobytes())
- elif isinstance(state, io.BytesIO):
- state.seek(0)
- return torch.load(state, map_location="cuda", weights_only=False)
- else:
- raise RuntimeError("Unsupported checkpoint format.")
-
- def _split_extra_state(self, state):
- fp8_checkpoint = self.fp8_meta["fp8_checkpoint"] or self.fp8 or self.fp8_calibration
-
- if not fp8_checkpoint:
- return [state] * self.num_gemms
-
- state = self._decode_extra_state(state)
- extra_states = []
- extra_fp8_variables = state['extra_fp8_variables']
- extra_fp8_variables['num_gemms'] = 1
- for gemm_idx in range(self.num_gemms):
- tmp_state = {"extra_fp8_variables": extra_fp8_variables}
- # TE 2.0 adds recipe in extra_state
- if is_te_min_version("2.0.0"):
- tmp_state['recipe'] = state['recipe']
- # Only delayed scaling has global fp8 meta tensors. We're not using
- # self.fp8_meta["recipe"].delayed() because it's available in TE 2.0 and later.
- if isinstance(self.fp8_meta["recipe"], te.common.recipe.DelayedScaling):
- tmp_state.update(
- {
- "scale_fwd": state['scale_fwd'].view(3, -1)[:, gemm_idx],
- "amax_history_fwd": state['amax_history_fwd'].view(
- self.fp8_meta["recipe"].amax_history_len, 3, -1
- )[:, :, gemm_idx],
- "scale_bwd": state['scale_bwd'].view(2, -1)[:, gemm_idx],
- "amax_history_bwd": state['amax_history_bwd'].view(
- self.fp8_meta["recipe"].amax_history_len, 2, -1
- )[:, :, gemm_idx],
- }
- )
- # TE 2.0 removes scale_inv_fwd and scale_inv_bwd
- if not is_te_min_version("2.0.0"):
- tmp_state.update(
- {
- "scale_inv_fwd": state['scale_inv_fwd'].view(3, -1)[:, gemm_idx],
- "scale_inv_bwd": state['scale_inv_bwd'].view(2, -1)[:, gemm_idx],
- }
- )
- extra_states.append(self._encode_extra_state(tmp_state))
- return extra_states
-
- def _sharded_state_dict_grouped(
- self, tp_axis_map, prefix='', sharded_offsets=(), metadata=None
- ):
- """
- prefix should be module_name to make keys identical to sequetial ones.
- """
- sharded_state_dict = {}
- full_state_dict = self.state_dict(prefix='', keep_vars=True)
- num_global_experts = get_expert_model_parallel_world_size() * self.num_gemms
- local_expert_indices_offset = get_expert_model_parallel_rank() * self.num_gemms
- ep_axis = len(sharded_offsets)
- extra_states = self._split_extra_state(full_state_dict['_extra_state'])
- for gemm_idx in range(self.num_gemms):
- state_dict = {
- f'{gemm_idx}.weight': full_state_dict[f'weight{gemm_idx}'],
- f'{gemm_idx}._extra_state': extra_states[gemm_idx],
- }
- if self.use_bias:
- state_dict[f'{gemm_idx}.bias'] = full_state_dict[f'bias{gemm_idx}']
- sub_sd = make_sharded_tensors_for_checkpoint(
- state_dict,
- '',
- tp_axis_map,
- (
- *sharded_offsets,
- (ep_axis, local_expert_indices_offset + gemm_idx, num_global_experts),
- ),
- )
- # Remove expert layers indexing from sharded keys
- replace_prefix_for_sharding(sub_sd, f'{gemm_idx}.', prefix)
- sharded_state_dict.update(
- {
- f'{prefix}weight{gemm_idx}': sub_sd[f'{gemm_idx}.weight'],
- f'{prefix}_extra_state{"" if gemm_idx == 0 else gemm_idx}': sub_sd[
- f'{gemm_idx}._extra_state'
- ],
- }
- )
- if self.use_bias:
- sharded_state_dict[f'{prefix}bias{gemm_idx}'] = sub_sd[f'{gemm_idx}.bias']
- # Adjust replica ids - replication along DP modulo EP
- for k, sh_ten in sharded_state_dict.items():
- replica_id = sh_ten.replica_id
- assert (
- len(replica_id) == 3
- ), f'Expected replica_id for {k} to be in (PP, TP, DP) format, got: {replica_id}'
- if getattr(sh_ten, "is_data_parallel_fully_shard", False):
- edp_replica_id = 0
- else:
- edp_replica_id = get_expert_data_parallel_rank()
- sh_ten.replica_id = (*replica_id[:2], edp_replica_id)
- return sharded_state_dict
-
- def backward_dw(self):
- """Compute weight gradients during the backward pass if split_bw is enabled."""
- if self.config.split_bw:
- super().backward_dw()
-
- class TEColumnParallelGroupedLinear(TEGroupedLinear):
- """
- Wrapper for the Transformer-Engine's `GroupedLinear` layer but specialized
- to column-parallel style.
- """
+ # softmax, dropout (if any), then matmul with v
+ attn = torch.softmax(scores, dim=-1)
+ if self.attn_dropout is not None:
+ attn = self.attn_dropout(attn)
- def __init__(
- self,
- num_gemms: int,
- input_size: int,
- output_size: int,
- *,
- config: ModelParallelConfig,
- init_method: Callable,
- bias: bool,
- skip_bias_add: bool,
- is_expert: bool,
- tp_comm_buffer_name: Optional[str] = None,
- ):
-
- super().__init__(
- num_gemms=num_gemms,
- input_size=input_size,
- output_size=output_size,
- parallel_mode="column",
- config=config,
- init_method=condition_init_method(config, init_method),
- bias=bias,
- skip_bias_add=skip_bias_add,
- is_expert=is_expert,
- tp_comm_buffer_name=tp_comm_buffer_name,
- )
+ # (B, H, Sq, Sk) @ (B, H, Sk, Dh) -> (B, H, Sq, Dh)
+ outBH = torch.matmul(attn, vBH)
- def sharded_state_dict(self, prefix='', sharded_offsets=(), metadata=None):
- """
- For each gemm, sharding along axis 0, bias sharded.
- Assume sharded_offsets[-1] is the expert parallel offset.
- """
- tp_axis_map = {}
- for gemm_idx in range(self.num_gemms):
- tp_axis_map.update({f'{gemm_idx}.weight': 0, f'{gemm_idx}.bias': 0})
- return super()._sharded_state_dict_grouped(
- tp_axis_map, prefix, sharded_offsets, metadata
- )
+ # bring back to (B, Sq, H, Dh)
+ out4 = outBH.transpose(1, 2).contiguous()
- class TERowParallelGroupedLinear(TEGroupedLinear):
- """
- Wrapper for the Transformer-Engine's `GroupedLinear` layer but specialized
- to row-parallel style.
- """
+ # If inputs were 3D originally, merge heads back
+ out = self._reconstruct_output(out4, q_was_3d)
- def __init__(
- self,
- num_gemms: int,
- input_size: int,
- output_size: int,
- *,
- config: ModelParallelConfig,
- init_method: Callable,
- bias: bool,
- skip_bias_add: bool,
- is_expert: bool,
- tp_comm_buffer_name: Optional[str] = None,
- ):
-
- super().__init__(
- num_gemms=num_gemms,
- input_size=input_size,
- output_size=output_size,
- parallel_mode="row",
- config=config,
- init_method=condition_init_method(config, init_method),
- bias=bias,
- skip_bias_add=skip_bias_add,
- is_expert=is_expert,
- tp_comm_buffer_name=tp_comm_buffer_name,
- )
+ return out
- def sharded_state_dict(self, prefix='', sharded_offsets=(), metadata=None):
- """
- For each gemm, sharding along axis 1, bias not sharded.
- Assume sharded_offsets[-1] is the expert parallel offset.
- """
- tp_axis_map = {f'{gemm_idx}.weight': 1 for gemm_idx in range(self.num_gemms)}
- return super()._sharded_state_dict_grouped(
- tp_axis_map, prefix, sharded_offsets, metadata
- )
+# ---------------------------------------------------------------------------
+# FP8 / RNG / misc TE utilities โ disabled
+# ---------------------------------------------------------------------------
+
+class TEDelayedScaling:
+ """
+ Stub for TE DelayedScaling FP8 recipe.
+ """
-else:
- TEGroupedLinear = None # type: ignore[assignment, misc]
- TEColumnParallelGroupedLinear = None # type: ignore[assignment, misc]
- TERowParallelGroupedLinear = None # type: ignore[assignment, misc]
+ def __init__(self, *args, **kwargs):
+ raise RuntimeError("FP8 / TEDelayedScaling is not available (Transformer Engine disabled).")
-class TEDelayedScaling(te.common.recipe.DelayedScaling):
+class TECudaRNGStatesTracker:
"""
- Wrapper for the Transformer-Engine's `DelayedScaling` layer.
+ Stub for TE CUDA RNG tracker.
+
+ Megatron's own RNG tracker should be used instead.
"""
- def __init__(
- self,
- config: ModelParallelConfig,
- fp8_format: int,
- override_linear_precision: tuple = (False, False, False),
- ):
- extra_kwargs = _get_extra_te_kwargs(config)
- if is_te_min_version("1.6.0.dev0"):
- extra_kwargs["fp8_dpa"] = config.fp8_dot_product_attention
- extra_kwargs["fp8_mha"] = config.fp8_multi_head_attention
- if get_te_version() < PkgVersion("1.8.0"):
- extra_kwargs["interval"] = config.fp8_interval
- elif config.fp8_interval != 1:
- warnings.warn("fp8_interval is deprecated and ignored from Transformer-Engine v1.8.0.")
+ def __init__(self, *args, **kwargs):
+ raise RuntimeError("TECudaRNGStatesTracker is not available (Transformer Engine disabled).")
- super().__init__(
- margin=config.fp8_margin,
- fp8_format=fp8_format,
- amax_compute_algo=config.fp8_amax_compute_algo,
- amax_history_len=config.fp8_amax_history_len,
- override_linear_precision=override_linear_precision,
- **extra_kwargs,
- )
+# ---------------------------------------------------------------------------
+# SplitAlongDim, RoPE, CPU offload โ minimal / disabled
+# ---------------------------------------------------------------------------
-class TECudaRNGStatesTracker(te.pytorch.distributed.CudaRNGStatesTracker):
- """Wraps TransformerEngine's CudaRNGStatesTracker so that it is
- interchangeable with Megatron's RNG tracker"""
+# def SplitAlongDim(x: Tensor, dim: int, num_chunks: int):
+# """
+# Minimal replacement for TE's _SplitAlongDim.apply.
- def __init__(self):
- super().__init__()
- self.reset()
+# Returns a tuple of chunks along the given dimension.
+# Autograd works fine through torch.chunk, so this is enough
+# for most use cases that only expect a "split" function.
+# """
+# return torch.chunk(x, num_chunks, dim=dim)
+import torch
+from torch import Tensor
+from typing import Union, Sequence
+
+def SplitAlongDim(x: Tensor, dim: int, num_chunks: Union[int, Sequence[int]]):
+ """
+ Replacement for TE's _SplitAlongDim.apply.
- def is_initialized(self):
- """Checks if the internal RNG state has been set wirth set_states()."""
- return self._is_initialized
+ - If `num_chunks` is an int: behaves like torch.chunk(x, num_chunks, dim=dim).
+ - If `num_chunks` is a sequence of ints: behaves like torch.split(x, num_chunks, dim=dim).
+ In that case the ints specify sizes for each split along `dim`.
- def reset(self):
- """Reset the internal RNG state."""
- super().reset()
- self._is_initialized = False
+ Returns a tuple of tensors.
- def set_states(self, states):
- """Set the internal RNG state."""
- super().set_states(states)
- self._is_initialized = True
+ Raises ValueError for invalid inputs (e.g. sizes don't sum to x.size(dim)).
+ """
+ if x is None:
+ raise ValueError("SplitAlongDim: input tensor x is None")
+
+ if not isinstance(dim, int):
+ raise TypeError(f"SplitAlongDim: dim must be int, got {type(dim)}")
+
+ # int -> torch.chunk
+ if isinstance(num_chunks, int):
+ if num_chunks <= 0:
+ raise ValueError(f"SplitAlongDim: num_chunks must be > 0, got {num_chunks}")
+ return tuple(torch.chunk(x, num_chunks, dim=dim))
+
+ # sequence -> torch.split with sizes
+ if isinstance(num_chunks, (list, tuple)):
+ sizes = list(num_chunks)
+ if not all(isinstance(s, int) and s >= 0 for s in sizes):
+ raise TypeError("SplitAlongDim: when passing sizes, all elements must be non-negative ints")
+ total = sum(sizes)
+ dim_size = x.size(dim)
+ if total != dim_size:
+ raise ValueError(
+ f"SplitAlongDim: sizes sum to {total} but tensor.size({dim}) == {dim_size}. "
+ "Sizes must match the dimension length exactly."
+ )
+ return tuple(torch.split(x, sizes, dim=dim))
- def add(self, name, seed):
- """Track the rng state."""
- super().add(name, seed)
- self._is_initialized = True
+ raise TypeError(
+ "SplitAlongDim: num_chunks must be an int (number of equal chunks) or "
+ "a list/tuple of sizes for each chunk."
+ )
-def te_checkpoint(
- forward_func,
- distribute_saved_activations,
- get_rng_state_tracker,
- tp_group,
- hidden_states,
- attention_mask,
- attn_mask_type,
- context,
- context_mask,
- rotary_pos_emb,
- **kwargs,
+def get_cpu_offload_context(
+ enabled: bool,
+ num_layers: int,
+ model_layers: Any,
+ activation_offloading: bool,
+ weight_offloading: bool,
):
- """Checkpointing with Transformer-Engine."""
- from transformer_engine.pytorch.distributed import checkpoint
-
- if is_te_min_version("1.5.0"):
- return checkpoint(
- forward_func,
- hidden_states,
- attention_mask,
- attn_mask_type,
- context,
- context_mask,
- rotary_pos_emb,
- distribute_saved_activations=distribute_saved_activations,
- get_rng_state_tracker=get_rng_state_tracker,
- tp_group=tp_group,
- **kwargs,
- )
- else:
- return checkpoint(
- forward_func,
- distribute_saved_activations,
- get_rng_state_tracker,
- tp_group,
- hidden_states,
- attention_mask,
- context,
- context_mask,
- rotary_pos_emb,
- )
+ """
+ Stub for TE CPU offload context.
+ Return a no-op context manager so Megatron's transformer block can always
+ use the result in a `with` statement, even when Transformer Engine is not
+ installed and CPU offload is disabled.
+ """
+ def _sync():
+ return None
+
+ return nullcontext(), _sync
+
+
+# def fused_apply_rotary_pos_emb(
+# t: torch.Tensor, freqs: torch.Tensor, transpose_output_memory: bool = False
+# ) -> torch.Tensor:
+# """
+# Stub for fused RoPE (sbhd). We fall back to naive RoPE application:
+
+# This is a placeholder; if your config explicitly requires TE fused RoPE,
+# you should not be using this stub file.
+# """
+# # A very naive RoPE: assume last dimension is split into (cos, sin)
+# # This is intentionally simple and may NOT match all TE behaviours.
+# dim = t.shape[-1]
+# half = dim // 2
+# cos = freqs[..., :half]
+# sin = freqs[..., half:half * 2]
+# x1, x2 = t[..., :half], t[..., half:half * 2]
+# rot = torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
+# if transpose_output_memory:
+# # We ignore this just like TE stub
+# pass
+# return rot
+import torch
+from torch import Tensor
+from typing import Tuple
+
+def _align_and_broadcast_to(tensor: Tensor, target: Tensor, name: str) -> Tensor:
+ """
+ Try to make `tensor` broadcastable to `target` by:
+ - checking last-dim size matches target's last-dim,
+ - if tensor has fewer dims: prepend singleton dims,
+ - if tensor has more dims: try to squeeze leading singleton dims,
+ - finally expand to target.shape.
-try:
- from transformer_engine.pytorch.attention import _SplitAlongDim
- SplitAlongDim = _SplitAlongDim.apply
+ If impossible, raise a helpful ValueError showing shapes.
+ """
+ if tensor is None:
+ raise ValueError(f"{name} is None")
-except ImportError:
- SplitAlongDim = None
+ if tensor.size(-1) != target.size(-1):
+ raise ValueError(
+ f"{name} last-dim size ({tensor.size(-1)}) != target last-dim ({target.size(-1)}). "
+ "They must match (both equal to half = D//2)."
+ )
-try:
- from transformer_engine.pytorch.cpu_offload import get_cpu_offload_context as _get_cpu_offload_context
- def get_cpu_offload_context(
- enabled, num_layers, model_layers, activation_offloading, weight_offloading
- ):
- """Get CPU offload context and sync function."""
- if is_te_min_version("1.10.0.dev0"):
- context, sync_func = _get_cpu_offload_context(
- enabled, num_layers, model_layers, activation_offloading, weight_offloading
- )
+ t_ndim = tensor.ndim
+ tgt_ndim = target.ndim
+
+ # If tensor has fewer dims, prepend singleton dims
+ if t_ndim < tgt_ndim:
+ new_shape = (1,) * (tgt_ndim - t_ndim) + tensor.shape
+ tensor = tensor.reshape(new_shape)
+
+ # If tensor has more dims, try to squeeze leading singleton dims
+ elif t_ndim > tgt_ndim:
+ # Attempt to remove leading dims of size 1 until dims match or cannot.
+ # Only allow squeezing if those dims are 1.
+ diff = t_ndim - tgt_ndim
+ leading = tensor.shape[:diff]
+ if all(s == 1 for s in leading):
+ tensor = tensor.reshape(tensor.shape[diff:])
else:
- context, sync_func = _get_cpu_offload_context(
- enabled, num_layers, activation_offloading, weight_offloading
+ # can't safely squeeze; give helpful message
+ raise ValueError(
+ f"{name} has more leading dims ({tensor.shape}) than target ({target.shape}) "
+ "and they are not singleton. Cannot broadcast. You probably need to slice freqs "
+ "to the appropriate time-range (e.g. freqs[sequence_start:sequence_end])."
)
- return context, sync_func
+ # Now tensor.ndim == target.ndim and last dims already checked equal.
+ try:
+ tensor = tensor.expand(target.shape)
+ except Exception as e:
+ raise ValueError(
+ f"Failed to broadcast {name} with shape {tensor.shape} to target shape {target.shape}: {e}"
+ )
+ return tensor
+
+
+# def fused_apply_rotary_pos_emb(
+# t: Tensor, freqs: Tensor, transpose_output_memory: bool = False
+# ) -> Tensor:
+# """
+# Memory-efficient, broadcasting-robust RoPE fallback.
+
+# - t: Tensor with shape (..., D) where D is even (D = 2*half)
+# - freqs: Tensor with last dim either half (cos only) or 2*half (cos||sin)
+# Leading dims of freqs are flexible but must be broadcastable to t's leading dims
+# after appropriate simple reshapes (see _align_and_broadcast_to).
+# - transpose_output_memory ignored (kept for API compatibility).
+# """
+# if t is None:
+# raise ValueError("fused_apply_rotary_pos_emb: input t is None")
+
+# D = t.size(-1)
+# if D % 2 != 0:
+# raise ValueError(f"fused_apply_rotary_pos_emb: last dimension must be even, got {D}")
+# half = D // 2
+
+# if freqs is None:
+# raise ValueError("fused_apply_rotary_pos_emb: freqs must not be None")
+
+# if freqs.size(-1) == 2 * half:
+# cos = freqs[..., :half]
+# sin = freqs[..., half : 2 * half]
+# elif freqs.size(-1) == half:
+# cos = freqs
+# sin = torch.zeros_like(cos)
+# else:
+# raise ValueError(
+# "fused_apply_rotary_pos_emb: freqs last dim must be half or 2*half "
+# f"(expected {half} or {2*half}, got {freqs.size(-1)})"
+# )
+
+# # Ensure cos/sin dtype/device match to avoid implicit copies in subsequent ops
+# if cos.dtype != t.dtype:
+# cos = cos.to(dtype=t.dtype)
+# if sin.dtype != t.dtype:
+# sin = sin.to(dtype=t.dtype)
+# if cos.device != t.device:
+# cos = cos.to(device=t.device)
+# if sin.device != t.device:
+# sin = sin.to(device=t.device)
+
+# # Extract x1, x2 views (no copy)
+# x1 = t[..., :half]
+# x2 = t[..., half : 2 * half]
+
+# # Align cos/sin to x1/x2 shapes (leading dims)
+# try:
+# cos_b = _align_and_broadcast_to(cos, x1, "cos")
+# sin_b = _align_and_broadcast_to(sin, x1, "sin")
+# except ValueError as e:
+# # augment message with shapes for debugging
+# raise ValueError(
+# f"fused_apply_rotary_pos_emb: cannot align freqs with inputs.\n"
+# f"t.shape={tuple(t.shape)}, x1.shape={tuple(x1.shape)}, freqs.shape={tuple(freqs.shape)}\n"
+# f"Detailed error: {e}"
+# )
+
+# # Allocate output once
+# out = torch.empty_like(t)
+
+# # Compute first half: out[..., :half] = x1 * cos - x2 * sin
+# # Use out= to reduce temporaries, then addcmul_ for the subtraction
+# torch.mul(x1, cos_b, out=out[..., :half]) # out[:half] = x1 * cos
+# out[..., :half].addcmul_(x2, -sin_b) # out[:half] += x2 * (-sin) -> x1*cos - x2*sin
+
+# # Compute second half: out[..., half:] = x1 * sin + x2 * cos
+# torch.mul(x1, sin_b, out=out[..., half:]) # out[half:] = x1 * sin
+# out[..., half:].addcmul_(x2, cos_b) # out[half:] += x2 * cos -> x1*sin + x2*cos
+
+# return out
+import torch
+from torch import Tensor
+from typing import Optional
-except ImportError:
- get_cpu_offload_context = None # type: ignore[assignment, misc]
+def _expand_freqs_for_packed_tokens(freqs: Tensor, cu_seqlens: Tensor, total_tokens: int) -> Tensor:
+ """
+ Expand freqs (shape [L, ...]) into packed layout according to cu_seqlens.
+ Returns shape (total_tokens, ...).
+ """
+ if cu_seqlens is None:
+ raise ValueError("_expand_freqs_for_packed_tokens: cu_seqlens is None")
+
+ if cu_seqlens.ndim != 1:
+ raise ValueError("_expand_freqs_for_packed_tokens: cu_seqlens must be 1D cumulative sums")
+
+ # Ensure using CPU for iteration (cheap)
+ cu = cu_seqlens.to("cpu")
+ seq_lengths = (cu[1:] - cu[:-1]).tolist()
+ if sum(seq_lengths) != total_tokens:
+ raise ValueError(
+ f"_expand_freqs_for_packed_tokens: sum(seq_lengths)={sum(seq_lengths)} != total_tokens={total_tokens}. "
+ "Check cu_seqlens and packed tensor construction."
+ )
-try:
- from transformer_engine.pytorch.attention import FusedRoPEFunc
- def fused_apply_rotary_pos_emb(
- t: torch.Tensor, freqs: torch.Tensor, transpose_output_memory: bool = False
- ) -> torch.Tensor:
- """Apply rotary positional embedding to input tensor T in `sbhd` format."""
- if transpose_output_memory:
- warnings.warn(
- "transpose_output_memory is not supported by TE's fused RoPE and will be ignored."
+ L = freqs.shape[0]
+ out_parts = []
+ for length in seq_lengths:
+ if length == 0:
+ continue
+ if length > L:
+ raise ValueError(
+ f"_expand_freqs_for_packed_tokens: requested length {length} > freqs available {L}. "
+ "You need freqs computed for this sequence length or slice differently."
)
- return FusedRoPEFunc.apply(t, freqs, "sbhd")
-
- def fused_apply_rotary_pos_emb_thd(
- t: torch.Tensor,
- cu_seqlens: torch.Tensor,
- freqs: torch.Tensor,
- cp_size: int = 1,
- cp_rank: int = 0,
- ) -> torch.Tensor:
- """
- Apply rotary positional embedding to input tensor T in `thd` format with CP support.
- """
- if is_te_min_version("1.11.0", check_equality=False):
- return FusedRoPEFunc.apply(t, freqs, "thd", cu_seqlens, cp_size, cp_rank)
+ out_parts.append(freqs[:length])
+ return torch.cat(out_parts, dim=0)
+
+
+def fused_apply_rotary_pos_emb(
+ t: Tensor,
+ freqs: Tensor,
+ transpose_output_memory: bool = False,
+ cu_seqlens: Optional[Tensor] = None,
+) -> Tensor:
+ """
+ Memory-efficient RoPE that supports packed sequences.
+
+ Args:
+ t: Tensor with shape (total_tokens, ... , D) where last dim D is even (D = 2*half).
+ Example from your code: t.shape == (6466, 32, 128).
+ freqs: Tensor with last dim either 2*half (cos||sin) or half (cos only).
+ Example: freqs.shape == (316,1,1,128) meaning 316 positions.
+ cu_seqlens: Optional 1D tensor length (batch+1) of cumulative sums that describe
+ how t is packed (used to expand freqs into a per-packed-token table).
+ Returns:
+ Tensor same shape as t with RoPE applied.
+ """
+
+ if t is None or freqs is None:
+ raise ValueError("fused_apply_rotary_pos_emb: t and freqs must not be None")
+
+ D = t.size(-1)
+ if D % 2 != 0:
+ raise ValueError(f"fused_apply_rotary_pos_emb: last dim must be even, got {D}")
+ half = D // 2
+
+ # extract cos/sin from freqs last dim
+ if freqs.size(-1) == 2 * half:
+ cos_all = freqs[..., :half]
+ sin_all = freqs[..., half : 2 * half]
+ elif freqs.size(-1) == half:
+ cos_all = freqs
+ sin_all = torch.zeros_like(cos_all)
+ else:
+ raise ValueError(
+ f"fused_apply_rotary_pos_emb: freqs last dim must be half or 2*half "
+ f"(expected {half} or {2*half}, got {freqs.size(-1)})"
+ )
+
+ total_tokens = t.shape[0]
+
+ # Expand freqs into per-packed-token table if cu_seqlens provided
+ if cu_seqlens is not None:
+ cos = _expand_freqs_for_packed_tokens(cos_all, cu_seqlens, total_tokens)
+ sin = _expand_freqs_for_packed_tokens(sin_all, cu_seqlens, total_tokens)
+ else:
+ # Try a few reasonable auto-alignments; otherwise raise informative error
+ if cos_all.shape[0] == total_tokens:
+ cos = cos_all
+ sin = sin_all
+ elif cos_all.shape[0] > 1 and total_tokens % cos_all.shape[0] == 0:
+ factor = total_tokens // cos_all.shape[0]
+ cos = cos_all.repeat_interleave(factor, dim=0)
+ sin = sin_all.repeat_interleave(factor, dim=0)
else:
- return FusedRoPEFunc.apply(t, freqs, "thd", cu_seqlens)
+ raise ValueError(
+ f"fused_apply_rotary_pos_emb: freqs time dim ({cos_all.shape[0]}) does not match t time dim ({total_tokens}). "
+ "Provide cu_seqlens to expand freqs to packed layout or slice freqs upstream."
+ )
-except ImportError:
- pass
+ # Move cos/sin to same dtype/device as t to avoid unexpected copies during ops
+ if cos.dtype != t.dtype:
+ cos = cos.to(dtype=t.dtype)
+ if sin.dtype != t.dtype:
+ sin = sin.to(dtype=t.dtype)
+ if cos.device != t.device:
+ cos = cos.to(device=t.device)
+ if sin.device != t.device:
+ sin = sin.to(device=t.device)
+
+ # Split t into x1, x2 (views)
+ x1 = t[..., :half]
+ x2 = t[..., half : 2 * half]
+
+ # Ensure cos/sin broadcast shape matches x1: add singleton dims if necessary
+ cos_b = cos
+ sin_b = sin
+ while cos_b.ndim < x1.ndim:
+ cos_b = cos_b.unsqueeze(1)
+ sin_b = sin_b.unsqueeze(1)
+
+ try:
+ cos_b = cos_b.expand(x1.shape)
+ sin_b = sin_b.expand(x1.shape)
+ except Exception as e:
+ raise ValueError(
+ f"fused_apply_rotary_pos_emb: failed to broadcast cos/sin to x1 shape.\n"
+ f"t.shape={tuple(t.shape)}, cos.shape={tuple(cos.shape)}, sin.shape={tuple(sin.shape)}\n"
+ f"error: {e}"
+ )
-try:
- from transformer_engine.pytorch import Fp8Padding, Fp8Unpadding # pylint: disable=unused-import
-except ImportError:
- Fp8Padding = None
- Fp8Unpadding = None
-
-try:
- from transformer_engine.pytorch.permutation import (
- moe_permute,
- moe_sort_chunks_by_index,
- moe_unpermute,
- )
- fused_permute = moe_permute
- fused_unpermute = moe_unpermute
- fused_sort_chunks_by_index = moe_sort_chunks_by_index
-
-except ImportError:
- fused_permute = None
- fused_unpermute = None
- fused_sort_chunks_by_index = None
+ # Compute into a single output tensor using in-place operations to minimize temporaries
+ out = torch.empty_like(t)
+ torch.mul(x1, cos_b, out=out[..., :half])
+ out[..., :half].addcmul_(x2, -sin_b)
+ torch.mul(x1, sin_b, out=out[..., half:])
+ out[..., half:].addcmul_(x2, cos_b)
+
+ return out
+
+
+# def fused_apply_rotary_pos_emb_thd(
+# t: torch.Tensor,
+# cu_seqlens: torch.Tensor,
+# freqs: torch.Tensor,
+# cp_size: int = 1,
+# cp_rank: int = 0,
+# ) -> torch.Tensor:
+# """
+# Stub for TE thd RoPE; we just delegate to the naive RoPE above and
+# ignore cu_seqlens / cp_*.
+# """
+# return fused_apply_rotary_pos_emb(t, freqs, transpose_output_memory=False)
+def fused_apply_rotary_pos_emb_thd(
+ t: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ freqs: torch.Tensor,
+ cp_size: int = 1,
+ cp_rank: int = 0,
+) -> torch.Tensor:
+ pass
diff --git a/aiak_megatron/megatron/core/extensions/transformer_engine2.py b/aiak_megatron/megatron/core/extensions/transformer_engine2.py
new file mode 100644
index 00000000..0a7833cd
--- /dev/null
+++ b/aiak_megatron/megatron/core/extensions/transformer_engine2.py
@@ -0,0 +1,724 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+"""transformer engine utilities - PURE PYTORCH FALLBACK (no transformer_engine).
+
+This module mirrors the public API used elsewhere (TENorm, TELinear, TEColumnParallelLinear,
+TERowParallelLinear, TEDotProductAttention, TEGroupedLinear, etc.) but implements them
+with pure PyTorch fallbacks so Transformer-Engine is not required.
+"""
+
+import dataclasses
+import io
+import os
+import pickle
+import warnings
+from contextlib import nullcontext
+from typing import Any, Callable, Optional, Tuple, Dict
+
+import torch
+import torch.nn as nn
+import math
+from packaging.version import Version as PkgVersion
+from torch import Tensor
+from torch.nn.parameter import Parameter
+
+# Megatron imports retained for compatibility (these should exist in your environment)
+from megatron.core.dist_checkpointing.utils import replace_prefix_for_sharding
+from megatron.core.model_parallel_config import ModelParallelConfig
+from megatron.core.packed_seq_params import PackedSeqParams
+from megatron.core.parallel_state import (
+ get_context_parallel_global_ranks,
+ get_context_parallel_group,
+ get_expert_data_parallel_rank,
+ get_expert_model_parallel_rank,
+ get_expert_model_parallel_world_size,
+ get_expert_tensor_parallel_group,
+ get_expert_tensor_parallel_rank,
+ get_expert_tensor_parallel_world_size,
+ get_hierarchical_context_parallel_groups,
+ get_tensor_model_parallel_group,
+ get_tensor_model_parallel_rank,
+ get_tensor_model_parallel_world_size,
+)
+from megatron.core.tensor_parallel import get_cuda_rng_tracker, get_expert_parallel_rng_tracker_name
+from megatron.core.tensor_parallel.layers import (
+ _initialize_affine_weight_cpu,
+ set_tensor_model_parallel_attributes,
+)
+from megatron.core.tensor_parallel.random import get_data_parallel_rng_tracker_name
+from megatron.core.tensor_parallel.utils import divide
+from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.enums import Fp8Recipe
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.transformer.utils import make_sharded_tensors_for_checkpoint
+from megatron.core.utils import get_te_version, is_te_min_version # these helpers are still useful
+
+# ------------------------
+# Small helpers & fallbacks
+# ------------------------
+
+def _get_extra_te_kwargs(config: TransformerConfig):
+ """Return minimal kwargs to mimic transformer-engine behavior. For fallback this is minimal."""
+ # Provide dtype and device guidance like TE would
+ params_dtype = getattr(config, "params_dtype", torch.float32)
+ extra = {"params_dtype": params_dtype}
+ # Expose device selection hints for CPU/meta initialization flags
+ if getattr(config, "use_cpu_initialization", False):
+ extra["device"] = "cpu"
+ elif getattr(config, "init_model_with_meta_device", False):
+ extra["device"] = "meta"
+ else:
+ # default to current cuda device if available
+ if torch.cuda.is_available():
+ extra["device"] = torch.cuda.current_device()
+ else:
+ extra["device"] = "cpu"
+ return extra
+
+
+def condition_init_method(config, init_method):
+ """Condition TE init_method on config.perform_initialization (mirror TE wrapper)."""
+ return init_method if getattr(config, "perform_initialization", True) else (lambda w: None)
+
+
+class _RMSNorm(nn.Module):
+ """Simple RMSNorm implementation used by fallback TENorm."""
+ def __init__(self, hidden_size: int, eps: float = 1e-8, zero_centered_gamma: bool = False, use_bias: bool = False):
+ super().__init__()
+ self.hidden_size = hidden_size
+ self.eps = eps
+ self.use_bias = use_bias
+ self.weight = nn.Parameter(torch.zeros(hidden_size) if zero_centered_gamma else torch.ones(hidden_size))
+ if use_bias:
+ self.bias = nn.Parameter(torch.zeros(hidden_size))
+ else:
+ self.register_parameter("bias", None)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ rms = x.pow(2).mean(dim=-1, keepdim=True).add(self.eps).sqrt()
+ x_norm = x / rms
+ out = x_norm * self.weight
+ if self.use_bias:
+ out = out + self.bias
+ return out
+
+
+class TENorm:
+ """Factory that returns LayerNorm or RMSNorm fallback matching TE behavior."""
+ def __new__(cls, config: TransformerConfig, hidden_size: int, eps: float = 1e-5):
+ zero_centered = getattr(config, "layernorm_zero_centered_gamma", False)
+ normalization = getattr(config, "normalization", "LayerNorm")
+ if normalization == "LayerNorm":
+ ln = nn.LayerNorm(normalized_shape=hidden_size, eps=eps, elementwise_affine=True)
+ # match TE zero-centered gamma behavior
+ with torch.no_grad():
+ ln.weight.fill_(0.0 if zero_centered else 1.0)
+ ln.bias.zero_()
+ return ln
+ elif normalization == "RMSNorm":
+ return _RMSNorm(hidden_size=hidden_size, eps=eps, zero_centered_gamma=zero_centered, use_bias=False)
+ else:
+ raise ValueError("Only LayerNorm and RMSNorm are currently supported in TENorm fallback.")
+
+
+def condition_init_method(config, init_method: Callable) -> Callable:
+ """Return provided init_method or fallback initializer (compat helper)."""
+ if callable(init_method):
+ return init_method
+ def _f(tensor: torch.Tensor):
+ nn.init.kaiming_uniform_(tensor, a=math.sqrt(5))
+ return _f
+
+
+# ------------------------
+# TELinear pure-PyTorch fallback
+# ------------------------
+
+def _apply_linear_with_multiplicity(x_flat: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor], in_features: int, out_features: int, te_return_bias: bool):
+ """
+ Robust linear application that handles last_dim equal to in_features or
+ last_dim divisible by in_features (treat as M blocks, apply linear to each block and average).
+ Returns (out, bias_or_none) where out has same leading dims as x_flat except last -> out_features.
+ """
+ last_dim = x_flat.shape[-1]
+ if last_dim == in_features:
+ out = torch.nn.functional.linear(x_flat, weight, bias)
+ return out, (bias if te_return_bias and bias is not None else None)
+ if last_dim % in_features == 0:
+ M = last_dim // in_features
+ lead = x_flat.shape[:-1]
+ # reshape: (..., M, in_features)
+ x_blocks = x_flat.reshape(*lead, M, in_features)
+ prod_lead = 1
+ for s in lead:
+ prod_lead *= s
+ x_blocks_flat = x_blocks.reshape(prod_lead * M, in_features)
+ out_blocks_flat = torch.nn.functional.linear(x_blocks_flat, weight, bias)
+ out_blocks = out_blocks_flat.view(*lead, M, out_features)
+ out = out_blocks.mean(dim=-2)
+ warnings.warn(
+ f"_apply_linear_with_multiplicity: last_dim {last_dim} interpreted as {M} blocks of {in_features}; "
+ "applied linear to each and averaged results.",
+ UserWarning,
+ )
+ return out, (bias if te_return_bias and bias is not None else None)
+ raise RuntimeError(
+ f"TELinear/_apply_linear_with_multiplicity: last_dim ({last_dim}) is not equal to in_features ({in_features}) "
+ "and is not divisible by in_features."
+ )
+
+
+class TELinear(nn.Module):
+ """Pure-PyTorch replacement implementing the TE Linear wrapper API surface."""
+ def __init__(
+ self,
+ input_size: int,
+ output_size: int,
+ *,
+ parallel_mode: Optional[str],
+ config: ModelParallelConfig,
+ init_method: Callable,
+ bias: bool,
+ skip_bias_add: bool,
+ skip_weight_param_allocation: bool,
+ tp_comm_buffer_name: Optional[str] = None,
+ is_expert: bool = False,
+ ):
+ super().__init__()
+ self.config = config
+ self.parallel_mode = parallel_mode
+ self.is_expert = is_expert
+
+ if skip_weight_param_allocation:
+ raise ValueError("skip_weight_param_allocation unsupported in TELinear fallback")
+
+ self.te_return_bias = skip_bias_add and bias
+ self.is_first_microbatch = True
+ self.disable_parameter_transpose_cache = getattr(config, "disable_parameter_transpose_cache", False)
+
+ if getattr(config, "sequence_parallel", False):
+ warnings.warn("sequence_parallel requested but unsupported in TELinear fallback; ignored.")
+
+ if self.parallel_mode not in (None, "duplicated", "column", "row"):
+ warnings.warn(f"Unknown parallel_mode={self.parallel_mode}; proceeding without TP semantics")
+
+ self.in_features = input_size
+ self.out_features = output_size
+
+ self.weight = nn.Parameter(torch.empty((self.out_features, self.in_features), dtype=getattr(config, "params_dtype", torch.float32)))
+ if bias:
+ self.bias = nn.Parameter(torch.empty(self.out_features, dtype=getattr(config, "params_dtype", torch.float32)))
+ else:
+ self.register_parameter("bias", None)
+
+ init_fn = condition_init_method(config, init_method)
+ try:
+ init_fn(self.weight)
+ except Exception:
+ nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+ if getattr(self, "bias", None) is not None:
+ nn.init.zeros_(self.bias)
+
+ for p in (p for p in (self.weight, getattr(self, "bias", None)) if p is not None):
+ if is_expert:
+ expert_parallel_size = getattr(self.config, "expert_model_parallel_size", 1)
+ setattr(p, "allreduce", not (expert_parallel_size > 1))
+ else:
+ setattr(p, "allreduce", True)
+ if parallel_mode == "duplicated":
+ setattr(p, "sequence_parallel", getattr(self.config, "sequence_parallel", False))
+
+ def set_tensor_parallel_group(self, tp_group):
+ warnings.warn("set_tensor_parallel_group called โ no TP support in TELinear fallback", UserWarning)
+ self.tp_group = tp_group
+
+ def forward(self, x: torch.Tensor):
+ if x is None:
+ raise RuntimeError("TELinear.forward received x=None")
+
+ orig_shape = tuple(x.shape)
+
+ def to_batch_first(t: torch.Tensor):
+ if t.ndim < 2:
+ return t, False
+ if t.shape[0] > t.shape[1]:
+ return t.transpose(0, 1), True
+ return t, False
+
+ if x.ndim == 4:
+ xb, was_seq_first = to_batch_first(x)
+ B, S, H, D = xb.shape
+ x_flat = xb.reshape(B, S, H * D)
+ restore_permute = was_seq_first
+ lead = (B, S)
+ elif x.ndim == 3:
+ xb, was_seq_first = to_batch_first(x)
+ B, S, HP = xb.shape
+ x_flat = xb
+ restore_permute = was_seq_first
+ lead = (B, S)
+ elif x.ndim == 2:
+ x_flat = x
+ restore_permute = False
+ lead = (x_flat.shape[0],)
+ else:
+ x_flat = x.contiguous().view(-1, x.shape[-1])
+ restore_permute = False
+ lead = x.shape[:-1]
+
+ out, returned_bias = _apply_linear_with_multiplicity(x_flat, self.weight, getattr(self, "bias", None), self.in_features, self.out_features, self.te_return_bias)
+
+ # restore layout
+ if x.ndim == 4:
+ if out.ndim == 2 and out.shape[0] == (lead[0] * lead[1]):
+ out = out.view(lead[0], lead[1], self.out_features)
+ if restore_permute:
+ out = out.transpose(0, 1)
+ elif x.ndim == 3:
+ if out.ndim == 2 and out.shape[0] == (lead[0] * lead[1]):
+ out = out.view(lead[0], lead[1], self.out_features)
+ if restore_permute:
+ out = out.transpose(0, 1)
+ elif x.ndim == 2:
+ pass
+ else:
+ out = out.view(*lead, self.out_features)
+
+ if self.te_return_bias:
+ return out, returned_bias
+ return out, None
+
+ def sharded_state_dict(self, prefix: str = "", sharded_offsets=(), metadata=None):
+ assert self.parallel_mode in (None, "duplicated"), "sharded_state_dict only supported for duplicated in fallback"
+ return self.state_dict(prefix=prefix)
+
+ def backward_dw(self):
+ if getattr(self.config, "split_bw", False):
+ warnings.warn("split_bw True but not implemented in TELinear fallback; no-op")
+
+
+# ------------------------
+# TELayerNormColumnParallelLinear fallback
+# ------------------------
+
+class TELayerNormColumnParallelLinear(nn.Module):
+ """Fallback for LayerNorm + Column-Parallel Linear combination (no TP)."""
+ def __init__(
+ self,
+ input_size: int,
+ output_size: int,
+ *,
+ config: TransformerConfig,
+ init_method: Callable,
+ gather_output: bool,
+ bias: bool,
+ skip_bias_add: bool,
+ is_expert: bool,
+ skip_weight_param_allocation: bool = False,
+ tp_comm_buffer_name: Optional[str] = None,
+ ):
+ super().__init__()
+ self.config = config
+
+ if gather_output:
+ raise ValueError("gather_output=True unsupported in fallback")
+ if is_expert:
+ raise ValueError("MoE (is_expert=True) unsupported in fallback")
+ if skip_weight_param_allocation:
+ raise ValueError("skip_weight_param_allocation unsupported in fallback")
+
+ self.te_return_bias = skip_bias_add and bias
+ self.is_first_microbatch = True
+ self.disable_parameter_transpose_cache = getattr(config, "disable_parameter_transpose_cache", False)
+
+ self.in_features = input_size
+ self.out_features = output_size
+ self.use_bias = bias
+
+ normalization = getattr(config, "normalization", "LayerNorm")
+ eps = getattr(config, "layernorm_epsilon", 1e-5)
+ zero_centered = getattr(config, "layernorm_zero_centered_gamma", False)
+ if normalization == "LayerNorm":
+ self.layernorm = nn.LayerNorm(normalized_shape=input_size, eps=eps, elementwise_affine=True)
+ if zero_centered:
+ with torch.no_grad():
+ self.layernorm.weight.zero_()
+ self.layernorm.bias.zero_()
+ elif normalization == "RMSNorm":
+ self.layernorm = _RMSNorm(input_size, eps=eps, zero_centered_gamma=zero_centered, use_bias=False)
+ else:
+ raise ValueError(f"Unsupported normalization: {normalization}")
+
+ self.weight = Parameter(torch.empty((self.out_features, self.in_features), dtype=getattr(config, "params_dtype", torch.float32)))
+ if self.use_bias:
+ self.bias = Parameter(torch.empty(self.out_features, dtype=getattr(config, "params_dtype", torch.float32)))
+ else:
+ self.register_parameter("bias", None)
+
+ init_fn = condition_init_method(config, init_method)
+ try:
+ init_fn(self.weight)
+ except Exception:
+ nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+ if getattr(self, "bias", None) is not None:
+ nn.init.zeros_(self.bias)
+
+ for p in (self.weight, self.bias) if self.bias is not None else (self.weight,):
+ setattr(p, "allreduce", True)
+ setattr(p, "sequence_parallel", getattr(config, "sequence_parallel", False))
+
+ def forward(self, x: torch.Tensor) -> Tuple[Tensor, Optional[Tensor]]:
+ _is_first_microbatch = None if self.disable_parameter_transpose_cache else self.is_first_microbatch
+ ln_out = self.layernorm(x)
+
+ # Flatten to apply linear helper robustly
+ lead = ln_out.shape[:-1]
+ flat = ln_out.contiguous().view(-1, ln_out.shape[-1])
+ out_flat, returned_bias = _apply_linear_with_multiplicity(flat, self.weight, getattr(self, "bias", None), self.in_features, self.out_features, self.te_return_bias)
+ out = out_flat.view(*lead, self.out_features)
+
+ self.is_first_microbatch = False
+ if self.te_return_bias:
+ return out, returned_bias
+ return out, None
+
+ def sharded_state_dict(self, prefix: str = "", sharded_offsets=(), metadata=None):
+ if not (torch.distributed.is_available() and torch.distributed.is_initialized()):
+ return self.state_dict(prefix=prefix)
+ world_size = torch.distributed.get_world_size()
+ rank = torch.distributed.get_rank()
+ out = self.out_features
+ per_shard = divide(out, world_size)
+ start = rank * per_shard
+ end = start + per_shard if rank != (world_size - 1) else out
+ sd = {}
+ sd[prefix + "weight"] = self.weight.data[start:end].clone()
+ if self.use_bias and "bias" in self._parameters:
+ sd[prefix + "bias"] = self.bias.data[start:end].clone()
+ return sd
+
+ def __repr__(self):
+ return f"{type(self).__name__}(in_features={self.in_features}, out_features={self.out_features}, bias={self.use_bias})"
+
+ def backward_dw(self):
+ if getattr(self.config, "split_bw", False):
+ warnings.warn("split_bw True but not implemented in fallback; no-op")
+
+
+# Column/Row wrappers (reuse TELinear)
+class TEColumnParallelLinear(TELinear):
+ def __init__(self, input_size: int, output_size: int, *, config: ModelParallelConfig, init_method: Callable, gather_output: bool, bias: bool, skip_bias_add: bool, is_expert: bool, skip_weight_param_allocation: bool = False, tp_comm_buffer_name: Optional[str] = None):
+ if gather_output:
+ raise ValueError('gather_output True unsupported in fallback')
+ super().__init__(
+ input_size=input_size,
+ output_size=output_size,
+ parallel_mode="column",
+ config=config,
+ init_method=(condition_init_method(config, init_method) if not getattr(config, "use_cpu_initialization", False) else (lambda w: None)),
+ bias=bias,
+ skip_bias_add=skip_bias_add,
+ skip_weight_param_allocation=skip_weight_param_allocation,
+ tp_comm_buffer_name=tp_comm_buffer_name,
+ is_expert=is_expert,
+ )
+
+
+class TERowParallelLinear(TELinear):
+ def __init__(self, input_size: int, output_size: int, *, config: ModelParallelConfig, init_method: Callable, bias: bool, input_is_parallel: bool, skip_bias_add: bool, is_expert: bool, tp_comm_buffer_name: Optional[str] = None):
+ if not input_is_parallel:
+ raise ValueError("input_is_parallel=False unsupported in fallback")
+ super().__init__(
+ input_size=input_size,
+ output_size=output_size,
+ parallel_mode="row",
+ config=config,
+ init_method=(condition_init_method(config, init_method) if not getattr(config, "use_cpu_initialization", False) else (lambda w: None)),
+ bias=bias,
+ skip_bias_add=skip_bias_add,
+ skip_weight_param_allocation=False,
+ tp_comm_buffer_name=tp_comm_buffer_name,
+ is_expert=is_expert,
+ )
+
+
+# ------------------------
+# Chunked attention helper + TEDotProductAttention fallback
+# ------------------------
+
+def sdp_chunked_attention(
+ q: torch.Tensor,
+ k: torch.Tensor,
+ v: torch.Tensor,
+ attention_mask: Optional[torch.Tensor],
+ dropout_p: float,
+ is_causal: bool,
+ use_sdp: bool,
+ mha_module: Optional[nn.MultiheadAttention] = None,
+ chunk_size: int = 64,
+):
+ """
+ Compute attention in chunks along the sequence axis to reduce peak memory.
+
+ q,k,v: (B, S, H, D)
+ attention_mask: SDP-compatible or None
+ chunk_size: number of query tokens per chunk
+ """
+ B, S, H, D = q.shape
+ out = torch.empty_like(q)
+
+ if use_sdp:
+ k_flat = k.reshape(B, S, H * D)
+ v_flat = v.reshape(B, S, H * D)
+ else:
+ embed_dim = H * D
+ k_flat = k.reshape(B, S, embed_dim)
+ v_flat = v.reshape(B, S, embed_dim)
+
+ have_full_mask = attention_mask is not None and attention_mask.ndim >= 2
+
+ for start in range(0, S, chunk_size):
+ end = min(start + chunk_size, S)
+ q_chunk = q[:, start:end] # (B, chunk, H, D)
+ if use_sdp:
+ q_flat = q_chunk.reshape(B, end - start, H * D)
+ if have_full_mask:
+ try:
+ attn_mask_chunk = attention_mask[:, start:end, :]
+ except Exception:
+ attn_mask_chunk = None
+ else:
+ attn_mask_chunk = None
+
+ out_flat_chunk = torch.nn.functional.scaled_dot_product_attention(
+ q_flat, k_flat, v_flat, attn_mask=attn_mask_chunk, dropout_p=dropout_p, is_causal=is_causal
+ )
+ out[:, start:end] = out_flat_chunk.reshape(B, end - start, H, D)
+ del q_flat, attn_mask_chunk, out_flat_chunk
+ torch.cuda.empty_cache()
+ else:
+ q_flat = q_chunk.reshape(B, end - start, H * D)
+ if attention_mask is not None:
+ warnings.warn("attention_mask ignored in chunked MHA fallback. Consider using SDP for masks.")
+ out_flat_chunk, _ = mha_module(q_flat, k_flat, v_flat, attn_mask=None)
+ out[:, start:end] = out_flat_chunk.reshape(B, end - start, H, D)
+ del q_flat, out_flat_chunk
+ torch.cuda.empty_cache()
+
+ return out
+
+
+class TEDotProductAttention(nn.Module):
+ """
+ Chunked, robust fallback attention wrapper. Accepts (B,S,H,D), (B,S,E), (S,B,H,D), (S,B,E).
+ """
+ cp_stream: Optional[torch.cuda.Stream] = None
+
+ def __init__(
+ self,
+ config: TransformerConfig,
+ layer_number: int,
+ attn_mask_type: AttnMaskType,
+ attention_type: str,
+ attention_dropout: Optional[float] = None,
+ softmax_scale: Optional[float] = None,
+ k_channels: Optional[int] = None,
+ v_channels: Optional[int] = None,
+ cp_comm_type: str = "p2p",
+ ):
+ super().__init__()
+ self.config = config
+ self.layer_number = layer_number
+ self.attn_mask_type = attn_mask_type
+ self.attention_type = attention_type
+ self.attention_dropout = attention_dropout if attention_dropout is not None else getattr(config, "attention_dropout", 0.0)
+ self.softmax_scale = softmax_scale
+ self.kv_channels = (k_channels, v_channels) if (k_channels is not None and v_channels is not None) else getattr(config, "kv_channels", None)
+ self.num_heads = getattr(config, "num_attention_heads_per_partition", None) or getattr(config, "num_attention_heads", None)
+ self.qkv_format = 'sbhd'
+ self.te_forward_mask_type = False
+
+ self.kept_packed_seq_params = set(field.name for field in dataclasses.fields(PackedSeqParams))
+ self._use_pytorch_sdp = hasattr(torch.nn.functional, "scaled_dot_product_attention")
+ self._mha: Optional[nn.MultiheadAttention] = None
+
+ # chunk size tuneable from config
+ self.attn_chunk_size = getattr(config, "attn_chunk_size", 64)
+
+ def _to_batch_first(self, t: torch.Tensor) -> Tuple[torch.Tensor, bool]:
+ if t.ndim < 3:
+ raise ValueError(f"Unsupported tensor ndim {t.ndim} in attention; expected >=3")
+ if t.shape[0] > t.shape[1]:
+ return t.transpose(0, 1), True
+ return t, False
+
+ def _ensure_b_s_h_d(self, tensor: torch.Tensor, name: str) -> Tuple[torch.Tensor, int, int, int, int]:
+ t_bsf, _ = self._to_batch_first(tensor)
+ if t_bsf.ndim == 4:
+ B, S, H, D = t_bsf.shape
+ return t_bsf, B, S, H, D
+ if t_bsf.ndim == 3:
+ if self.num_heads is None:
+ raise RuntimeError(f"num_attention_heads unknown; cannot reshape 3D {name} tensor")
+ B, S, E = t_bsf.shape
+ H = self.num_heads
+ if E % H != 0:
+ raise ValueError(
+ f"Embedding dim {E} (for {name}) is not divisible by num_heads {H}."
+ )
+ D = E // H
+ t_bsf = t_bsf.view(B, S, H, D)
+ return t_bsf, B, S, H, D
+ raise ValueError(f"Unsupported tensor ndim for {name}: {tensor.ndim}")
+
+ def forward(
+ self,
+ query: Tensor,
+ key: Tensor,
+ value: Tensor,
+ attention_mask: Optional[Tensor],
+ attn_mask_type: AttnMaskType,
+ attention_bias: Optional[Tensor] = None,
+ packed_seq_params: Optional[PackedSeqParams] = None,
+ ) -> Tensor:
+ q, B, S, H, D = self._ensure_b_s_h_d(query, "query")
+ k, _, _, _, _ = self._ensure_b_s_h_d(key, "key")
+ v, _, _, _, _ = self._ensure_b_s_h_d(value, "value")
+
+ if q.numel() == 0 or k.numel() == 0 or v.numel() == 0:
+ raise RuntimeError("Empty q/k/v passed to attention")
+
+ is_causal = attn_mask_type == AttnMaskType.causal
+
+ if not self._use_pytorch_sdp and self._mha is None:
+ if self.num_heads is None:
+ raise RuntimeError("num_attention_heads unknown for MultiheadAttention fallback")
+ embed_dim = H * D
+ self._mha = nn.MultiheadAttention(embed_dim=embed_dim, num_heads=self.num_heads, dropout=self.attention_dropout, batch_first=True)
+
+ chunk_size = max(1, int(getattr(self, "attn_chunk_size", 64)))
+ out = sdp_chunked_attention(
+ q, k, v,
+ attention_mask,
+ dropout_p=self.attention_dropout,
+ is_causal=is_causal,
+ use_sdp=self._use_pytorch_sdp,
+ mha_module=self._mha,
+ chunk_size=chunk_size,
+ )
+ return out
+
+
+# ------------------------
+# Grouped linear fallbacks
+# ------------------------
+
+class TEGroupedLinear(nn.Module):
+ """Grouped linear fallback implemented as ModuleList of TELinear."""
+ def __init__(self, num_gemms: int, input_size: int, output_size: int, *, parallel_mode: Optional[str], config: ModelParallelConfig, init_method: Callable, bias: bool, skip_bias_add: bool, is_expert: bool = False, tp_comm_buffer_name: Optional[str] = None):
+ super().__init__()
+ self.num_gemms = num_gemms
+ self.config = config
+ self.submodules = nn.ModuleList(
+ [
+ TELinear(input_size, output_size, parallel_mode=parallel_mode, config=config, init_method=init_method, bias=bias, skip_bias_add=skip_bias_add, skip_weight_param_allocation=False, tp_comm_buffer_name=tp_comm_buffer_name, is_expert=is_expert)
+ for _ in range(num_gemms)
+ ]
+ )
+ self.te_return_bias = skip_bias_add and bias
+ self.is_first_microbatch = True
+ self.disable_parameter_transpose_cache = getattr(config, "disable_parameter_transpose_cache", False)
+
+ def forward(self, x, m_splits=None):
+ outs = []
+ for module in self.submodules:
+ out = module(x)
+ outs.append(out[0] if isinstance(out, tuple) else out)
+ stacked = torch.stack(outs, dim=0)
+ if self.te_return_bias:
+ biases = [m.bias for m in self.submodules if getattr(m, "bias", None) is not None]
+ return stacked, biases
+ return stacked, None
+
+
+class TEColumnParallelGroupedLinear(TEGroupedLinear):
+ def __init__(self, num_gemms: int, input_size: int, output_size: int, *, config: ModelParallelConfig, init_method: Callable, bias: bool, skip_bias_add: bool, is_expert: bool, tp_comm_buffer_name: Optional[str] = None):
+ super().__init__(num_gemms=num_gemms, input_size=input_size, output_size=output_size, parallel_mode="column", config=config, init_method=condition_init_method(config, init_method), bias=bias, skip_bias_add=skip_bias_add, is_expert=is_expert, tp_comm_buffer_name=tp_comm_buffer_name)
+
+
+class TERowParallelGroupedLinear(TEGroupedLinear):
+ def __init__(self, num_gemms: int, input_size: int, output_size: int, *, config: ModelParallelConfig, init_method: Callable, bias: bool, skip_bias_add: bool, is_expert: bool, tp_comm_buffer_name: Optional[str] = None):
+ super().__init__(num_gemms=num_gemms, input_size=input_size, output_size=output_size, parallel_mode="row", config=config, init_method=condition_init_method(config, init_method), bias=bias, skip_bias_add=skip_bias_add, is_expert=is_expert, tp_comm_buffer_name=tp_comm_buffer_name)
+
+
+# ------------------------
+# Stubs and TE-like utilities
+# ------------------------
+
+class TEDelayedScaling:
+ """Stub for TE DelayedScaling; FP8 not implemented in fallback."""
+ def __init__(self, config: ModelParallelConfig, fp8_format: int, override_linear_precision: tuple = (False, False, False)):
+ warnings.warn("TEDelayedScaling stub: FP8 not implemented in pure-PyTorch fallback.", UserWarning)
+ self.config = config
+
+
+class TECudaRNGStatesTracker:
+ """Small wrapper mimicking TE's RNG tracker surface."""
+ def __init__(self):
+ self._is_initialized = False
+ def reset(self):
+ self._is_initialized = False
+ def is_initialized(self):
+ return self._is_initialized
+ def set_states(self, states):
+ self._is_initialized = True
+ def add(self, name, seed):
+ self._is_initialized = True
+
+
+def fused_apply_rotary_pos_emb(t: torch.Tensor, freqs: torch.Tensor, transpose_output_memory: bool = False) -> torch.Tensor:
+ """Fallback non-fused RoPE. This is a simple placeholder; keep if your code expects the function to exist."""
+ warnings.warn("Using non-fused RoPE fallback. Consider replacing with optimized implementation.", UserWarning)
+ # naive elementwise application for sbhd layout: t: (..., H, D) and freqs shaped accordingly.
+ # This fallback simply returns input unchanged (safe default). If you need actual RoPE behavior,
+ # provide your project's RoPE implementation here.
+ return t
+
+
+def fused_apply_rotary_pos_emb_thd(t: torch.Tensor, cu_seqlens: torch.Tensor, freqs: torch.Tensor, cp_size: int = 1, cp_rank: int = 0) -> torch.Tensor:
+ warnings.warn("Using non-fused RoPE-thd fallback. Consider replacing with optimized implementation.", UserWarning)
+ return t
+
+
+fused_permute = None
+fused_unpermute = None
+fused_sort_chunks_by_index = None
+
+Fp8Padding = None
+Fp8Unpadding = None
+
+
+def get_cpu_offload_context(enabled, num_layers, model_layers, activation_offloading, weight_offloading):
+ warnings.warn("get_cpu_offload_context fallback: no TE CPU-offload context available.", UserWarning)
+ return nullcontext(), (lambda: None)
+
+
+def te_checkpoint(
+ forward_func,
+ distribute_saved_activations,
+ get_rng_state_tracker,
+ tp_group,
+ hidden_states,
+ attention_mask,
+ attn_mask_type,
+ context,
+ context_mask,
+ rotary_pos_emb,
+ **kwargs,
+):
+ """Fallback checkpointing using torch.utils.checkpoint (best-effort)."""
+ import torch.utils.checkpoint as cp
+ # Best-effort wrapper: TE's checkpoint signature differs across versions; we use a simple approach
+ return cp.checkpoint(forward_func, hidden_states, attention_mask, context)
+
+
+# End of pure-PyTorch fallback module
diff --git a/aiak_megatron/megatron/core/extensions/transformer_engine_org.py b/aiak_megatron/megatron/core/extensions/transformer_engine_org.py
new file mode 100644
index 00000000..d4c8cd61
--- /dev/null
+++ b/aiak_megatron/megatron/core/extensions/transformer_engine_org.py
@@ -0,0 +1,1452 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+"""transformer engine utilities"""
+
+import dataclasses
+import io
+import os
+import pickle
+import warnings
+from typing import Any, Callable, Optional
+
+import torch
+import transformer_engine as te
+from packaging.version import Version as PkgVersion
+from torch import Tensor
+from torch.nn.parameter import Parameter
+
+from megatron.core.dist_checkpointing.utils import replace_prefix_for_sharding
+from megatron.core.model_parallel_config import ModelParallelConfig
+from megatron.core.packed_seq_params import PackedSeqParams
+from megatron.core.parallel_state import (
+ get_context_parallel_global_ranks,
+ get_context_parallel_group,
+ get_expert_data_parallel_rank,
+ get_expert_model_parallel_rank,
+ get_expert_model_parallel_world_size,
+ get_expert_tensor_parallel_group,
+ get_expert_tensor_parallel_rank,
+ get_expert_tensor_parallel_world_size,
+ get_hierarchical_context_parallel_groups,
+ get_tensor_model_parallel_group,
+ get_tensor_model_parallel_rank,
+ get_tensor_model_parallel_world_size,
+)
+from megatron.core.tensor_parallel import get_cuda_rng_tracker, get_expert_parallel_rng_tracker_name
+from megatron.core.tensor_parallel.layers import (
+ _initialize_affine_weight_cpu,
+ set_tensor_model_parallel_attributes,
+)
+from megatron.core.tensor_parallel.random import get_data_parallel_rng_tracker_name
+from megatron.core.tensor_parallel.utils import divide
+from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.enums import Fp8Recipe
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.transformer.utils import make_sharded_tensors_for_checkpoint
+from megatron.core.utils import get_te_version, is_te_min_version
+
+
+def _get_extra_te_kwargs(config: TransformerConfig):
+ extra_transformer_engine_kwargs = {"params_dtype": config.params_dtype}
+
+ if is_te_min_version("0.12.0"):
+ if config.use_cpu_initialization:
+ extra_transformer_engine_kwargs["device"] = 'cpu'
+ elif config.init_model_with_meta_device:
+ extra_transformer_engine_kwargs["device"] = "meta"
+ else:
+ extra_transformer_engine_kwargs["device"] = torch.cuda.current_device()
+ return extra_transformer_engine_kwargs
+
+
+def condition_init_method(config, init_method):
+ """Condition TE init_method on config.perform_initialization."""
+ return init_method if config.perform_initialization else (lambda w: None)
+
+
+class TENorm:
+ """
+ A conditional wrapper to initialize an instance of Transformer-Engine's
+ `LayerNorm` or `RMSNorm` based on input
+ """
+
+ # TODO should we ditch normalization config and just use spec to choose LayerNorm vs RMSNorm?
+ def __new__(cls, config: TransformerConfig, hidden_size: int, eps: float = 1e-5):
+ if config.normalization == "LayerNorm":
+ instance = te.pytorch.LayerNorm(
+ hidden_size=hidden_size,
+ eps=eps,
+ sequence_parallel=config.sequence_parallel,
+ zero_centered_gamma=config.layernorm_zero_centered_gamma,
+ **_get_extra_te_kwargs(config),
+ )
+ elif config.normalization == "RMSNorm":
+ assert hasattr(
+ te.pytorch, "RMSNorm"
+ ), "Transformer-Engine >= v0.11 required to use this feature"
+ instance = te.pytorch.RMSNorm(
+ hidden_size=hidden_size,
+ eps=eps,
+ sequence_parallel=config.sequence_parallel,
+ zero_centered_gamma=config.layernorm_zero_centered_gamma,
+ **_get_extra_te_kwargs(config),
+ )
+ else:
+ raise Exception('Only LayerNorm and RMSNorm are curently supported')
+
+ return instance
+
+
+class TELinear(te.pytorch.Linear):
+ """
+ Wrapper for the Transformer-Engine's `Linear` layer.
+
+ Note that if Megatron's parallel_state has not been initialized
+ yet, the tp_group passed to TE will be None and must be set later
+ via set_tensor_parallel_group().
+
+ parallel_mode currently supports 3 different values:
+ - "column": Split the weight matrix along output dimension (used in TEColumnParallelLinear)
+ - "row": Split the weight matrix along input dimension (used in TERowParallelLinear)
+ - "duplicated": No tensor parallelism and weight is duplicated across TP ranks
+ - Note: For expert linear layers, we will disable communication logic here
+ as TP communication is handled in token_dispatcher.
+ """
+
+ def __init__(
+ self,
+ input_size: int,
+ output_size: int,
+ *,
+ parallel_mode: Optional[str],
+ config: ModelParallelConfig,
+ init_method: Callable,
+ bias: bool,
+ skip_bias_add: bool,
+ skip_weight_param_allocation: bool,
+ tp_comm_buffer_name: Optional[str] = None,
+ is_expert: bool = False,
+ ):
+ self.config = config
+
+ # TE returns a zero length Tensor when bias=False and
+ # return_bias=True, but we prefer None. So in that case we
+ # tell TE to not return the bias, and return None
+ # ourselves. This way our forward always returns two values
+ # and we don't have to deal with the zero length Tensor.
+ self.te_return_bias = skip_bias_add and bias
+ self.is_first_microbatch = True
+ self.disable_parameter_transpose_cache = self.config.disable_parameter_transpose_cache
+ if skip_weight_param_allocation:
+ raise ValueError(
+ 'Transformer Engine linear layers do not support skip_weight_param_allocation'
+ )
+
+ extra_kwargs = _get_extra_te_kwargs(config)
+
+ if self.config.split_bw:
+ extra_kwargs["delay_wgrad_compute"] = self.config.split_bw
+
+ if (
+ self.config.tp_comm_overlap
+ and tp_comm_buffer_name
+ and tp_comm_buffer_name not in ['qkv', 'proj', 'fc1', 'fc2']
+ ):
+ self.config.tp_comm_overlap = False
+ warnings.warn(
+ f"The user buffer name {tp_comm_buffer_name} is not supported in"
+ "Transformer Engine. Disabling TP communication overlap "
+ "for this layer."
+ )
+
+ if is_te_min_version("0.8.0"):
+ if self.config.tp_comm_overlap:
+ if is_te_min_version("1.5.0"):
+ # Use old overlap flags if they were supplied instead
+ extra_kwargs["ub_overlap_ag"] = (
+ self.config.tp_comm_overlap_ag
+ if hasattr(self.config, "tp_comm_overlap_ag")
+ else self.config.tp_comm_split_ag or self.config.tp_comm_atomic_ag
+ )
+ extra_kwargs["ub_overlap_rs"] = (
+ self.config.tp_comm_overlap_rs
+ if hasattr(self.config, "tp_comm_overlap_rs")
+ else self.config.tp_comm_split_rs or self.config.tp_comm_atomic_rs
+ )
+ # Disable ub overlap for experts.
+ if is_expert:
+ extra_kwargs["ub_overlap_ag"] = False
+ extra_kwargs["ub_overlap_rs"] = False
+ else:
+ extra_kwargs["ub_split_ag"] = self.config.tp_comm_split_ag
+ extra_kwargs["ub_atomic_gemm_ag"] = self.config.tp_comm_atomic_ag
+ extra_kwargs["ub_split_rs"] = self.config.tp_comm_split_rs
+ extra_kwargs["ub_atomic_gemm_rs"] = self.config.tp_comm_atomic_rs
+ # Disable ub overlap for experts.
+ if is_expert:
+ extra_kwargs["ub_split_ag"] = False
+ extra_kwargs["ub_atomic_gemm_ag"] = False
+ extra_kwargs["ub_split_rs"] = False
+ extra_kwargs["ub_atomic_gemm_rs"] = False
+ if is_te_min_version("1.0.0", check_equality=False):
+ assert (
+ tp_comm_buffer_name is not None
+ ), "Buffer name should be set to configure communication overlap settings"
+ extra_kwargs["ub_name"] = tp_comm_buffer_name
+
+ self.expert_parallel = self.config.expert_model_parallel_size > 1
+ if is_expert:
+ rng_tracker_name = get_expert_parallel_rng_tracker_name()
+ else:
+ if parallel_mode == "duplicated":
+ rng_tracker_name = get_data_parallel_rng_tracker_name()
+ else:
+ rng_tracker_name = None
+ if is_te_min_version("1.7.0"):
+ extra_kwargs["rng_tracker_name"] = rng_tracker_name
+
+ te_parallel_mode = parallel_mode
+ if parallel_mode == "duplicated":
+ # Handle non-parallel case
+ tp_group = None
+ tp_size = 1
+ explicit_expert_comm = False
+ te_parallel_mode = None
+ else:
+ # Disable communications in TE when using TP or EP by
+ # making TE agnostic of model parallel.
+ if is_expert:
+ tp_group = get_expert_tensor_parallel_group(check_initialized=False)
+ tp_size = get_expert_tensor_parallel_world_size()
+ else:
+ tp_group = get_tensor_model_parallel_group(check_initialized=False)
+ tp_size = get_tensor_model_parallel_world_size()
+ explicit_expert_comm = is_expert and (tp_size > 1 or self.expert_parallel)
+
+ if explicit_expert_comm:
+ if parallel_mode == "column":
+ output_size = divide(output_size, tp_size)
+ elif parallel_mode == "row":
+ input_size = divide(input_size, tp_size)
+ te_parallel_mode = None
+ tp_size = 1
+ tp_group = None
+
+
+ super().__init__(
+ in_features=input_size,
+ out_features=output_size,
+ sequence_parallel=self.config.sequence_parallel,
+ fuse_wgrad_accumulation=self.config.gradient_accumulation_fusion,
+ tp_group=tp_group,
+ tp_size=tp_size,
+ get_rng_state_tracker=(
+ get_cuda_rng_tracker if get_cuda_rng_tracker().is_initialized() else None
+ ),
+ init_method=condition_init_method(config, init_method),
+ bias=bias,
+ return_bias=self.te_return_bias,
+ parallel_mode=te_parallel_mode,
+ **extra_kwargs,
+ )
+
+ for param in self.parameters():
+ if is_expert:
+ # Reduce the gradient on the expert_data_parallel group for expert linear layers
+ setattr(param, 'allreduce', not self.expert_parallel)
+ else:
+ # Reduce the gradient on DP group
+ setattr(param, 'allreduce', True)
+ if parallel_mode == "duplicated":
+ # Reduce the gradient further on the TP group since the weight is
+ # duplicated across TP ranks
+ setattr(param, 'sequence_parallel', self.config.sequence_parallel)
+
+ def forward(self, x):
+ """Forward."""
+ _is_first_microbatch = (
+ None if self.disable_parameter_transpose_cache else self.is_first_microbatch
+ )
+
+ out = super().forward(x, is_first_microbatch=_is_first_microbatch)
+ self.is_first_microbatch = False
+
+ # TE only returns a tuple when return_bias is True, otherwise
+ # it returns a single Tensor, we always want to return two
+ # values regardless of the arguments.
+ if self.te_return_bias:
+ return out
+ return out, None
+
+ def sharded_state_dict(self, prefix='', sharded_offsets=(), metadata=None):
+ """Replicate cross TP/DP."""
+
+ # Provide the dist-ckpt support when TELinear is directly used
+ # It can only happen with duplicated parallel mode
+ assert (
+ self.parallel_mode is None
+ ), "TELinear sharded_state_dict can only be used with duplicated parallel mode"
+ state_dict = self.state_dict(prefix='', keep_vars=True)
+ return make_sharded_tensors_for_checkpoint(state_dict, prefix, None, sharded_offsets)
+
+ def backward_dw(self):
+ """Compute weight gradients during the backward pass if split_bw is enabled."""
+ if self.config.split_bw:
+ super().backward_dw()
+
+class TELayerNormColumnParallelLinear(te.pytorch.LayerNormLinear):
+ """
+ Wrapper for the Transformer-Engine's `LayerNormLinear` layer that combines
+ layernorm and linear layers
+ """
+
+ def __init__(
+ self,
+ input_size: int,
+ output_size: int,
+ *,
+ config: TransformerConfig,
+ init_method: Callable,
+ gather_output: bool,
+ bias: bool,
+ skip_bias_add: bool,
+ is_expert: bool,
+ skip_weight_param_allocation: bool = False,
+ tp_comm_buffer_name: Optional[str] = None,
+ ):
+ self.config = config
+
+ if gather_output:
+ raise ValueError('Transformer Engine linear layers do not support gather_output = True')
+
+ if is_expert:
+ raise ValueError('Transformer Engine linear layers do not yet support MoE')
+
+ if skip_weight_param_allocation:
+ raise ValueError(
+ 'Transformer Engine linear layers do not support skip_weight_param_allocation'
+ )
+
+ # TE returns a zero length Tensor when bias=False and
+ # return_bias=True, but we prefer None. So in that case we
+ # tell TE to not return the bias, and return None
+ # ourselves. This way our forward always returns two values
+ # and we don't have to deal with the zero length Tensor.
+ self.te_return_bias = skip_bias_add and bias
+ self.is_first_microbatch = True
+ self.disable_parameter_transpose_cache = self.config.disable_parameter_transpose_cache
+ extra_kwargs = _get_extra_te_kwargs(config)
+
+ if self.config.split_bw:
+ extra_kwargs["delay_wgrad_compute"] = self.config.split_bw
+
+ # Only Transformer-Engine version >= 0.11.0 supports `RMSNorm`
+ if is_te_min_version("0.11.0"):
+ extra_kwargs["normalization"] = self.config.normalization
+ elif self.config.normalization != "LayerNorm":
+ te_version = get_te_version()
+ raise ValueError(
+ f"Transformer Engine v{te_version} does not support {self.config.normalization}."
+ )
+
+ if is_te_min_version("0.8.0"):
+ if self.config.tp_comm_overlap:
+ extra_kwargs["ub_bulk_wgrad"] = self.config.tp_comm_bulk_wgrad
+ extra_kwargs["ub_bulk_dgrad"] = self.config.tp_comm_bulk_dgrad
+ if is_te_min_version("1.5.0", check_equality=False):
+ # Use old overlap flags if they were supplied instead
+ extra_kwargs["ub_overlap_ag"] = (
+ self.config.tp_comm_overlap_ag
+ if hasattr(self.config, "tp_comm_overlap_ag")
+ else self.config.tp_comm_split_ag or self.config.tp_comm_atomic_ag
+ )
+ if is_te_min_version("1.6.0.dev0", check_equality=False):
+ extra_kwargs["ub_overlap_rs_dgrad"] = (
+ self.config.tp_comm_overlap_rs_dgrad
+ if hasattr(self.config, "tp_comm_overlap_rs_dgrad")
+ else False
+ )
+ if tp_comm_buffer_name == 'qkv' and self.config.tp_comm_overlap_disable_qkv:
+ extra_kwargs["ub_overlap_ag"] = False
+ extra_kwargs["ub_overlap_rs_dgrad"] = False
+
+ if tp_comm_buffer_name == 'fc1' and self.config.tp_comm_overlap_disable_fc1:
+ extra_kwargs["ub_overlap_ag"] = False
+ extra_kwargs["ub_overlap_rs_dgrad"] = False
+ else:
+ extra_kwargs["ub_atomic_gemm_ag"] = self.config.tp_comm_atomic_ag
+ extra_kwargs["ub_split_ag"] = self.config.tp_comm_split_ag
+ if is_te_min_version("1.0.0", check_equality=False):
+ assert (
+ tp_comm_buffer_name is not None
+ ), "Buffer name should be set to configure communication overlap settings"
+ extra_kwargs["ub_name"] = tp_comm_buffer_name
+
+
+ super().__init__(
+ in_features=input_size,
+ out_features=output_size,
+ eps=self.config.layernorm_epsilon,
+ sequence_parallel=self.config.sequence_parallel,
+ fuse_wgrad_accumulation=self.config.gradient_accumulation_fusion,
+ tp_group=get_tensor_model_parallel_group(check_initialized=False),
+ tp_size=self.config.tensor_model_parallel_size,
+ get_rng_state_tracker=(
+ get_cuda_rng_tracker if get_cuda_rng_tracker().is_initialized() else None
+ ),
+ init_method=(
+ condition_init_method(config, init_method)
+ if not config.use_cpu_initialization
+ else lambda w: None
+ ),
+ bias=bias,
+ return_bias=self.te_return_bias,
+ parallel_mode="column",
+ return_layernorm_output=False,
+ zero_centered_gamma=self.config.layernorm_zero_centered_gamma,
+ **extra_kwargs,
+ )
+
+ world_size = get_tensor_model_parallel_world_size()
+ rank = get_tensor_model_parallel_rank()
+
+ if config.use_cpu_initialization:
+ output_size_per_partition = divide(output_size, world_size)
+ _ = _initialize_affine_weight_cpu(
+ self.weight,
+ output_size,
+ input_size,
+ output_size_per_partition,
+ 0,
+ init_method=condition_init_method(config, init_method),
+ stride=1,
+ return_master_weight=False,
+ rank=rank,
+ world_size=world_size,
+ skip_set_tensor_parallel_attributes=True,
+ )
+ if bias:
+ self.bias = Parameter(
+ torch.empty(output_size_per_partition, dtype=config.params_dtype)
+ )
+ set_tensor_model_parallel_attributes(self.bias, True, 0, 1)
+ with torch.no_grad():
+ self.bias.zero_()
+ setattr(self.bias, 'allreduce', True)
+
+ def forward(self, x):
+ """Forward."""
+ _is_first_microbatch = (
+ None if self.disable_parameter_transpose_cache else self.is_first_microbatch
+ )
+
+ out = super().forward(x, is_first_microbatch=_is_first_microbatch)
+ self.is_first_microbatch = False
+
+ # TE only returns a tuple when return_bias is True, otherwise
+ # it returns a single Tensor, we always want to return two
+ # values regardless of the arguments.
+ if self.te_return_bias:
+ return out
+ return out, None
+
+ def sharded_state_dict(self, prefix='', sharded_offsets=(), metadata=None):
+ """Sharding along axis 0, bias sharded"""
+ state_dict = self.state_dict(prefix='', keep_vars=True)
+ return make_sharded_tensors_for_checkpoint(
+ state_dict, prefix, {'weight': 0, 'bias': 0}, sharded_offsets
+ )
+
+ def __repr__(self):
+ return (
+ f"{type(self).__name__}(in_features={self.in_features}, "
+ f"out_features={self.out_features}, bias={self.use_bias}, TP={self.tp_size})"
+ )
+
+ def backward_dw(self):
+ """Compute weight gradients during the backward pass if split_bw is enabled."""
+ if self.config.split_bw:
+ super().backward_dw()
+
+
+class TEColumnParallelLinear(TELinear):
+ """
+ Wrapper for the Transformer-Engine's `Linear` layer but specialized similar
+ to megatron's `ColumnParallelLinear` layer.
+ """
+
+ def __init__(
+ self,
+ input_size: int,
+ output_size: int,
+ *,
+ config: ModelParallelConfig,
+ init_method: Callable,
+ gather_output: bool,
+ bias: bool,
+ skip_bias_add: bool,
+ is_expert: bool,
+ skip_weight_param_allocation: bool = False,
+ tp_comm_buffer_name: Optional[str] = None,
+ ):
+ if gather_output:
+ raise ValueError('Transformer Engine linear layers do not support gather_output = True')
+
+ super().__init__(
+ input_size=input_size,
+ output_size=output_size,
+ parallel_mode="column",
+ config=config,
+ init_method=(
+ condition_init_method(config, init_method)
+ if not config.use_cpu_initialization
+ else lambda w: None
+ ),
+ bias=bias,
+ skip_bias_add=skip_bias_add,
+ is_expert=is_expert,
+ skip_weight_param_allocation=skip_weight_param_allocation,
+ tp_comm_buffer_name=tp_comm_buffer_name,
+ )
+
+ if config.use_cpu_initialization:
+ if is_expert:
+ world_size = get_expert_tensor_parallel_world_size()
+ rank = get_expert_tensor_parallel_rank()
+ else:
+ world_size = get_tensor_model_parallel_world_size()
+ rank = get_tensor_model_parallel_rank()
+ output_size_per_partition = divide(output_size, world_size)
+ _ = _initialize_affine_weight_cpu(
+ self.weight,
+ output_size,
+ input_size,
+ output_size_per_partition,
+ 0,
+ init_method=condition_init_method(config, init_method),
+ stride=1,
+ return_master_weight=False,
+ rank=rank,
+ world_size=world_size,
+ skip_set_tensor_parallel_attributes=True,
+ )
+ if bias:
+ self.bias = Parameter(
+ torch.empty(output_size_per_partition, dtype=config.params_dtype)
+ )
+ set_tensor_model_parallel_attributes(self.bias, True, 0, 1)
+ with torch.no_grad():
+ self.bias.zero_()
+ setattr(self.bias, 'allreduce', True)
+
+ def sharded_state_dict(self, prefix='', sharded_offsets=(), metadata=None):
+ """Sharding along axis 0, bias sharded"""
+ state_dict = self.state_dict(prefix='', keep_vars=True)
+ return make_sharded_tensors_for_checkpoint(
+ state_dict, prefix, {'weight': 0, 'bias': 0}, sharded_offsets
+ )
+
+ def __repr__(self):
+ return (
+ f"{type(self).__name__}(in_features={self.in_features}, "
+ f"out_features={self.out_features}, bias={self.use_bias}, TP={self.tp_size})"
+ )
+
+ def backward_dw(self):
+ """Compute weight gradients during the backward pass if split_bw is enabled."""
+ if self.config.split_bw:
+ super().backward_dw()
+
+
+class TERowParallelLinear(TELinear):
+ """
+ Wrapper for the Transformer-Engine's `Linear` layer but specialized similar
+ to megatron's `RowParallelLinear` layer.
+ """
+
+ def __init__(
+ self,
+ input_size: int,
+ output_size: int,
+ *,
+ config: ModelParallelConfig,
+ init_method: Callable,
+ bias: bool,
+ input_is_parallel: bool,
+ skip_bias_add: bool,
+ is_expert: bool,
+ tp_comm_buffer_name: Optional[str] = None,
+ ):
+ if not input_is_parallel:
+ raise ValueError(
+ "Transformer Engine linear layers do not support input_is_parallel = False"
+ )
+
+ super().__init__(
+ input_size=input_size,
+ output_size=output_size,
+ parallel_mode="row",
+ config=config,
+ init_method=(
+ condition_init_method(config, init_method)
+ if not config.use_cpu_initialization
+ else lambda w: None
+ ),
+ bias=bias,
+ skip_bias_add=skip_bias_add,
+ skip_weight_param_allocation=False, # We don't currently use this for row parallel layers
+ is_expert=is_expert,
+ tp_comm_buffer_name=tp_comm_buffer_name,
+ )
+ if config.use_cpu_initialization:
+ if is_expert:
+ world_size = get_expert_tensor_parallel_world_size()
+ rank = get_expert_tensor_parallel_rank()
+ else:
+ world_size = get_tensor_model_parallel_world_size()
+ rank = get_tensor_model_parallel_rank()
+ input_size_per_partition = divide(input_size, world_size)
+ self.master_weight = _initialize_affine_weight_cpu(
+ self.weight,
+ output_size,
+ input_size,
+ input_size_per_partition,
+ 1,
+ init_method=condition_init_method(config, init_method),
+ stride=1,
+ return_master_weight=False,
+ params_dtype=config.params_dtype,
+ rank=rank,
+ world_size=world_size,
+ skip_set_tensor_parallel_attributes=True,
+ )
+ if bias:
+ self.bias = Parameter(torch.empty(output_size, dtype=config.params_dtype))
+ # Always initialize bias to zero.
+ with torch.no_grad():
+ self.bias.zero_()
+ setattr(self.bias, 'allreduce', True)
+ setattr(self.bias, 'sequence_parallel', config.sequence_parallel)
+
+ def sharded_state_dict(self, prefix='', sharded_offsets=(), metadata=None):
+ """Sharding along axis 1, bias not sharded"""
+ state_dict = self.state_dict(prefix='', keep_vars=True)
+ return make_sharded_tensors_for_checkpoint(
+ state_dict, prefix, {'weight': 1}, sharded_offsets
+ )
+
+ def __repr__(self):
+ return (
+ f"{type(self).__name__}(in_features={self.in_features}, "
+ f"out_features={self.out_features}, bias={self.use_bias}, TP={self.tp_size})"
+ )
+
+ def backward_dw(self):
+ """Compute weight gradients during the backward pass if split_bw is enabled."""
+ if self.config.split_bw:
+ super().backward_dw()
+
+
+class TEDotProductAttention(te.pytorch.DotProductAttention):
+ """
+ Wrapper for the Transformer-Engine's `DotProductAttention` layer that also
+ has "flash attention" enabled.
+
+ Note that if Megatron's parallel_state has not been initialized yet, the
+ tp_group and cp_group passed to TE will be None and must be set later
+ via set_tensor_parallel_group() and set_context_parallel_group().
+ """
+
+ cp_stream: torch.cuda.Stream = None
+
+ def __init__(
+ self,
+ config: TransformerConfig,
+ layer_number: int,
+ attn_mask_type: AttnMaskType,
+ attention_type: str,
+ attention_dropout: Optional[float] = None,
+ softmax_scale: Optional[float] = None,
+ k_channels: Optional[int] = None,
+ v_channels: Optional[int] = None,
+ cp_comm_type: str = "p2p",
+ ):
+ self.config = config
+ self.te_forward_mask_type = False
+ self.qkv_format: str = 'sbhd'
+
+ if self.config.apply_query_key_layer_scaling != bool(
+ int(os.getenv('NVTE_APPLY_QK_LAYER_SCALING', '0'))
+ ):
+ raise ValueError(
+ f"apply_query_key_layer_scaling is {self.config.apply_query_key_layer_scaling} "
+ f"but environment variable NVTE_APPLY_QK_LAYER_SCALING is "
+ f"{os.getenv('NVTE_APPLY_QK_LAYER_SCALING')}. Transformer Engine does not support "
+ f"setting query key layer scaling via argument, so these two must match."
+ )
+
+ extra_kwargs: dict[str, Any] = {}
+ if is_te_min_version("0.11.0"):
+ extra_kwargs["num_gqa_groups"] = self.config.num_query_groups
+ elif self.config.num_query_groups != self.config.num_attention_heads:
+ raise ValueError(
+ f"Transformer Engine v{get_te_version()} does not support Grouped Query Attention, "
+ f"use a newer version of Transformer Engine. "
+ f"(num_query_groups ({self.config.num_query_groups}) != "
+ f"num_attention_heads ({self.config.num_attention_heads}))"
+ )
+
+ if is_te_min_version("0.10.0"):
+ extra_kwargs["attention_type"] = attention_type
+ # older version don't need attention_type
+
+ if is_te_min_version("0.12.0", check_equality=False):
+ self.te_forward_mask_type = True
+
+ # This check is important as CP config can be disabled while having a valid CP group
+ # Example - Disabling CP for encoder while a valid CP group exists for decoder
+ if self.config.context_parallel_size > 1:
+ assert is_te_min_version("1.0.0"), "Only Transformer-Engine version >= 1.0.0 supports context parallelism!"
+ if getattr(TEDotProductAttention, "cp_stream") is None:
+ TEDotProductAttention.cp_stream = torch.cuda.Stream()
+
+ extra_kwargs["cp_group"] = get_context_parallel_group(check_initialized=False)
+ extra_kwargs["cp_global_ranks"] = get_context_parallel_global_ranks(check_initialized=False)
+ extra_kwargs["cp_stream"] = TEDotProductAttention.cp_stream
+
+ if is_te_min_version("1.10.0"):
+ if cp_comm_type is None:
+ extra_kwargs["cp_comm_type"] = "p2p"
+
+ elif cp_comm_type == "a2a+p2p":
+ assert is_te_min_version("1.12.0"), (
+ f"Transformer-Engine v{get_te_version()} must be >= 1.12.0 to support"
+ "hierarchical cp commucation."
+ )
+ extra_kwargs["cp_comm_type"] = "a2a+p2p"
+ extra_kwargs["cp_group"] = get_hierarchical_context_parallel_groups(check_initialized=False)
+ else:
+ extra_kwargs["cp_comm_type"] = cp_comm_type
+
+ if self.config.deterministic_mode:
+ if int(os.getenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1")) != 0:
+ raise RuntimeError(
+ "deterministic_mode is on and we are using DotProductAttention from "
+ "Transformer Engine, but NVTE_ALLOW_NONDETERMINISTIC_ALGO is not 0. "
+ f"Currently set to: {os.getenv('NVTE_ALLOW_NONDETERMINISTIC_ALGO', 'not set')}."
+ )
+
+ if config.window_size is not None:
+ # Check version
+ assert is_te_min_version("1.2.0"), (
+ f"Transformer-Engine v{get_te_version()} must be >= 1.2.0 to support sliding window attention."
+ )
+ extra_kwargs['window_size'] = config.window_size
+
+ if is_te_min_version("1.10.0"):
+ # TE 1.10.0 introduces the ability to set the different k and v channels
+ kv_channels = (
+ (k_channels, v_channels)
+ if k_channels is not None and v_channels is not None
+ else self.config.kv_channels
+ )
+ extra_kwargs['softmax_scale'] = softmax_scale
+ else:
+ kv_channels = self.config.kv_channels
+
+ self.kept_packed_seq_params = set(field.name for field in dataclasses.fields(PackedSeqParams))
+ if get_te_version() < PkgVersion("1.3.0"):
+ # TE 1.3.0 introduces precomputing max_seqlen to remove unnecessary kernels and D2H
+ # copies (#555)
+ # These two arguments did not exist prior to 1.3.0
+ self.kept_packed_seq_params.discard("max_seqlen_q")
+ self.kept_packed_seq_params.discard("max_seqlen_kv")
+
+ if get_te_version() < PkgVersion("1.10.0"):
+ # TE 1.8.0 introduces cu_seqlens_padded which is the cu_seqlens with paddings counted
+ # in each individual sequence in THD format dataset
+ # These two arguments did not exist prior to 1.8.0. Full support added in 1.10.0 (#1012)
+ self.kept_packed_seq_params.discard("cu_seqlens_q_padded")
+ self.kept_packed_seq_params.discard("cu_seqlens_kv_padded")
+
+ super().__init__(
+ num_attention_heads=self.config.num_attention_heads,
+ kv_channels=kv_channels,
+ attention_dropout=self.config.attention_dropout if attention_dropout is None else attention_dropout,
+ attn_mask_type=attn_mask_type.name,
+ sequence_parallel=self.config.sequence_parallel,
+ tp_size=self.config.tensor_model_parallel_size,
+ get_rng_state_tracker=(
+ get_cuda_rng_tracker if get_cuda_rng_tracker().is_initialized() else None
+ ),
+ tp_group=get_tensor_model_parallel_group(check_initialized=False),
+ layer_number=layer_number,
+ **extra_kwargs,
+ )
+
+ def forward(
+ self,
+ query: Tensor,
+ key: Tensor,
+ value: Tensor,
+ attention_mask: Tensor,
+ attn_mask_type: AttnMaskType,
+ attention_bias: Tensor = None,
+ packed_seq_params: PackedSeqParams = None,
+ ):
+ """Forward."""
+ packed_seq_kwargs = (
+ {key: getattr(packed_seq_params, key) for key in self.kept_packed_seq_params}
+ if packed_seq_params is not None
+ else {}
+ )
+ # overwrite self.qkv_format depending on self.config.apply_rope_fusion, which can be set
+ # after init
+ if self.config.apply_rope_fusion and is_te_min_version("0.13.0", check_equality=False):
+ self.qkv_format = 'bshd'
+
+ qkv_format = packed_seq_kwargs.get('qkv_format', self.qkv_format)
+
+ # WAR for peak memory usage.
+ # See https://gitlab-master.nvidia.com/ADLR/megatron-lm/-/merge_requests/2388
+ if self.config.apply_rope_fusion and qkv_format == 'bshd':
+ query, key, value = [x.transpose(0, 1).contiguous() for x in (query, key, value)]
+ # In PyTorch, the following two tensors are in fact the same:
+ # Tensor with shape (1, S, H, D) and stride (S*H*D, H*D, D, 1)
+ # Tensor with shape (1, S, H, D) and stride (H*D, H*D, D, 1)
+ # Stride for a dimension that is 1 has no meaning, so tensors created two different ways
+ # can have same shape but different strides.
+ # We unify them to the first one to pass the stride check in TE
+ if value.shape == key.shape and value.shape[0] == 1 and value.stride() != key.stride():
+ value = value.as_strided(value.shape, key.stride())
+
+ attention_bias_kwargs = {}
+ if attention_bias is not None:
+ assert is_te_min_version("1.2.0"), (
+ f"Transformer-Engine v{get_te_version()} must be >= 1.2.0 to support `attention_bias`."
+ )
+ attention_bias_kwargs = dict(
+ core_attention_bias_type='post_scale_bias', core_attention_bias=attention_bias
+ )
+ else:
+ if self.config.position_embedding_type == "alibi":
+ attention_bias_kwargs = dict(core_attention_bias_type='alibi', core_attention_bias=None)
+
+ if self.te_forward_mask_type:
+ if qkv_format == 'thd' and is_te_min_version("1.7.0"):
+ # thd format uses flash attention with cuDNN kernel which requires is_padding=True,
+ # so the only acceptable mask types are `padding_causal` and `padding`. These do not
+ # necessarily indicate there are padded tokens in the sequence.
+ if attn_mask_type == AttnMaskType.causal:
+ attn_mask_type = AttnMaskType.padding_causal
+ elif attn_mask_type == AttnMaskType.no_mask:
+ attn_mask_type = AttnMaskType.padding
+ core_attn_out = super().forward(
+ query,
+ key,
+ value,
+ attention_mask,
+ attn_mask_type=attn_mask_type.name,
+ **attention_bias_kwargs,
+ **packed_seq_kwargs,
+ )
+ else:
+ core_attn_out = super().forward(
+ query,
+ key,
+ value,
+ attention_mask,
+ **attention_bias_kwargs,
+ **packed_seq_kwargs,
+ )
+
+ if self.config.apply_rope_fusion and qkv_format == 'bshd':
+ return core_attn_out.transpose(0, 1)
+ else:
+ return core_attn_out
+
+
+if is_te_min_version("1.9.0.dev0"):
+
+ class TEGroupedLinear(te.pytorch.GroupedLinear):
+ """
+ Wrapper for the Transformer-Engine's `GroupedLinear` layer.
+
+ Note that if Megatron's parallel_state has not been initialized
+ yet, the tp_group passed to TE will be None and must be set later
+ via set_tensor_parallel_group().
+ """
+
+ def __init__(
+ self,
+ num_gemms: int,
+ input_size: int,
+ output_size: int,
+ *,
+ parallel_mode: Optional[str],
+ config: ModelParallelConfig,
+ init_method: Callable,
+ bias: bool,
+ skip_bias_add: bool,
+ is_expert: bool = False,
+ tp_comm_buffer_name: Optional[str] = None,
+ ):
+ self.config = config
+
+ # TE returns a zero length Tensor when bias=False and
+ # return_bias=True, but we prefer None. So in that case we
+ # tell TE to not return the bias, and return None
+ # ourselves. This way our forward always returns two values
+ # and we don't have to deal with the zero length Tensor.
+ self.te_return_bias = skip_bias_add and bias
+ self.is_first_microbatch = True
+ self.disable_parameter_transpose_cache = self.config.disable_parameter_transpose_cache
+
+ extra_kwargs = _get_extra_te_kwargs(config)
+
+ if self.config.split_bw:
+ extra_kwargs["delay_wgrad_compute"] = self.config.split_bw
+
+ extra_kwargs["ub_name"] = tp_comm_buffer_name
+
+ self.expert_parallel = self.config.expert_model_parallel_size > 1
+ if is_expert:
+ extra_kwargs["rng_tracker_name"] = get_expert_parallel_rng_tracker_name()
+
+ # The comms between TP and EP group is explicitly handled by MoE token dispatcher.
+ # So we disable comms by making TE agnostic of model parallel.
+ if is_expert:
+ tp_group = get_expert_tensor_parallel_group(check_initialized=False)
+ tp_size = get_expert_tensor_parallel_world_size()
+ else:
+ tp_group = get_tensor_model_parallel_group(check_initialized=False)
+ tp_size = get_tensor_model_parallel_world_size()
+ self.explicit_expert_comm = is_expert and (tp_size > 1 or self.expert_parallel)
+
+ if self.explicit_expert_comm:
+ if parallel_mode == "column":
+ output_size = divide(output_size, tp_size)
+ elif parallel_mode == "row":
+ input_size = divide(input_size, tp_size)
+ parallel_mode = None
+ tp_size = 1
+ tp_group = None
+
+
+ super().__init__(
+ num_gemms=num_gemms,
+ in_features=input_size,
+ out_features=output_size,
+ sequence_parallel=self.config.sequence_parallel,
+ fuse_wgrad_accumulation=self.config.gradient_accumulation_fusion,
+ tp_group=tp_group,
+ tp_size=tp_size,
+ get_rng_state_tracker=(
+ get_cuda_rng_tracker if get_cuda_rng_tracker().is_initialized() else None
+ ),
+ init_method=condition_init_method(config, init_method),
+ bias=bias,
+ return_bias=self.te_return_bias,
+ parallel_mode=parallel_mode,
+ **extra_kwargs,
+ )
+
+ for param in self.parameters():
+ setattr(param, 'allreduce', not (is_expert and self.expert_parallel))
+
+ def merge_extra_states(
+ self,
+ state_dict,
+ prefix,
+ local_metadata,
+ strict,
+ missing_keys,
+ unexpected_keys,
+ error_msgs,
+ ):
+ """
+ Merge multiple "_extra_state" into one.
+ """
+ self.init_fp8_metadata(num_gemms=self.num_gemms)
+ # When resume training, loading ckpt is out of fp8_autocast context.
+ # So we need to manually detect from the state_dict.
+ fp8_checkpoint = any("_extra_state" in str(key) for key in state_dict.keys())
+
+ if not fp8_checkpoint:
+ return
+
+ try:
+ state_list = [
+ state_dict.pop(f"{prefix}_extra_state{i}") for i in range(1, self.num_gemms)
+ ]
+ except KeyError:
+ # "_extra_state{i}" only exists for dist-ckpt. Return for torch native ckpt.
+ return
+
+ # Early return conditions:
+ # 1. Empty state_dict
+ # 2. Empty state_list
+ # 3. _extra_state is None
+ # 4. _extra_state does not contain any information
+ if (
+ not state_dict
+ or not state_list
+ or state_dict.get(f"{prefix}_extra_state") is None
+ or self._decode_extra_state(state_dict[f"{prefix}_extra_state"]) is None
+ ):
+ return
+
+ state_list = [state_dict.pop(f"{prefix}_extra_state")] + state_list
+ state_list = [self._decode_extra_state(state) for state in state_list]
+ extra_fp8_variables = state_list[0]['extra_fp8_variables']
+ extra_fp8_variables['num_gemms'] = self.num_gemms
+ extra_state = {"extra_fp8_variables": extra_fp8_variables}
+ # TE 2.0 adds recipe in extra_state
+ if is_te_min_version("2.0.0"):
+ self.fp8_meta["recipe"] = state_list[0]['recipe']
+ extra_state['recipe'] = self.fp8_meta["recipe"]
+ # Only delayed scaling has global fp8 meta tensors. We're not using
+ # self.fp8_meta["recipe"].delayed() because it's available in TE 2.0 and later.
+ if isinstance(self.fp8_meta["recipe"], te.common.recipe.DelayedScaling):
+ extra_state.update(
+ {
+ "scale_fwd": torch.cat(
+ [state['scale_fwd'].view(-1, 1) for state in state_list], dim=1
+ ).view(-1),
+ "amax_history_fwd": torch.cat(
+ [state['amax_history_fwd'].view(-1, 1) for state in state_list],
+ dim=1,
+ ).view(self.fp8_meta["recipe"].amax_history_len, -1),
+ "scale_bwd": torch.cat(
+ [state['scale_bwd'].view(-1, 1) for state in state_list], dim=1
+ ).view(-1),
+ "amax_history_bwd": torch.cat(
+ [state['amax_history_bwd'].view(-1, 1) for state in state_list],
+ dim=1,
+ ).view(self.fp8_meta["recipe"].amax_history_len, -1),
+ }
+ )
+ # TE 2.0 removes scale_inv_fwd and scale_inv_bwd
+ if not is_te_min_version("2.0.0"):
+ extra_state.update(
+ {
+ "scale_inv_fwd": torch.cat(
+ [state['scale_inv_fwd'].view(-1, 1) for state in state_list],
+ dim=1,
+ ).view(-1),
+ "scale_inv_bwd": torch.cat(
+ [state['scale_inv_bwd'].view(-1, 1) for state in state_list],
+ dim=1,
+ ).view(-1),
+ }
+ )
+ state_dict[f"{prefix}_extra_state"] = self._encode_extra_state(extra_state)
+
+ self._register_load_state_dict_pre_hook(merge_extra_states, with_module=True)
+
+ def forward(self, x, m_splits):
+ """Forward."""
+ _is_first_microbatch = (
+ None if self.disable_parameter_transpose_cache else self.is_first_microbatch
+ )
+ out = super().forward(x, m_splits, is_first_microbatch=_is_first_microbatch)
+ self.is_first_microbatch = False
+
+ # TE only returns a tuple when return_bias is True, otherwise
+ # it returns a single Tensor, we always want to return two
+ # values regardless of the arguments.
+ if self.te_return_bias:
+ return out
+ return out, None
+
+ def _encode_extra_state(self, state):
+ # TE 2.0 changed the format of extra_state to be a byte tensor
+ if is_te_min_version("2.0.0"):
+ torch.cuda.synchronize()
+ state_serialized = bytearray(pickle.dumps(state))
+ state_serialized = torch.frombuffer(state_serialized, dtype=torch.uint8)
+ else:
+ state_serialized = io.BytesIO()
+ torch.save(state, state_serialized)
+ return state_serialized
+
+ def _decode_extra_state(self, state):
+ if isinstance(state, torch.Tensor):
+ return pickle.loads(state.detach().cpu().numpy().tobytes())
+ elif isinstance(state, io.BytesIO):
+ state.seek(0)
+ return torch.load(state, map_location="cuda", weights_only=False)
+ else:
+ raise RuntimeError("Unsupported checkpoint format.")
+
+ def _split_extra_state(self, state):
+ fp8_checkpoint = self.fp8_meta["fp8_checkpoint"] or self.fp8 or self.fp8_calibration
+
+ if not fp8_checkpoint:
+ return [state] * self.num_gemms
+
+ state = self._decode_extra_state(state)
+ extra_states = []
+ extra_fp8_variables = state['extra_fp8_variables']
+ extra_fp8_variables['num_gemms'] = 1
+ for gemm_idx in range(self.num_gemms):
+ tmp_state = {"extra_fp8_variables": extra_fp8_variables}
+ # TE 2.0 adds recipe in extra_state
+ if is_te_min_version("2.0.0"):
+ tmp_state['recipe'] = state['recipe']
+ # Only delayed scaling has global fp8 meta tensors. We're not using
+ # self.fp8_meta["recipe"].delayed() because it's available in TE 2.0 and later.
+ if isinstance(self.fp8_meta["recipe"], te.common.recipe.DelayedScaling):
+ tmp_state.update(
+ {
+ "scale_fwd": state['scale_fwd'].view(3, -1)[:, gemm_idx],
+ "amax_history_fwd": state['amax_history_fwd'].view(
+ self.fp8_meta["recipe"].amax_history_len, 3, -1
+ )[:, :, gemm_idx],
+ "scale_bwd": state['scale_bwd'].view(2, -1)[:, gemm_idx],
+ "amax_history_bwd": state['amax_history_bwd'].view(
+ self.fp8_meta["recipe"].amax_history_len, 2, -1
+ )[:, :, gemm_idx],
+ }
+ )
+ # TE 2.0 removes scale_inv_fwd and scale_inv_bwd
+ if not is_te_min_version("2.0.0"):
+ tmp_state.update(
+ {
+ "scale_inv_fwd": state['scale_inv_fwd'].view(3, -1)[:, gemm_idx],
+ "scale_inv_bwd": state['scale_inv_bwd'].view(2, -1)[:, gemm_idx],
+ }
+ )
+ extra_states.append(self._encode_extra_state(tmp_state))
+ return extra_states
+
+ def _sharded_state_dict_grouped(
+ self, tp_axis_map, prefix='', sharded_offsets=(), metadata=None
+ ):
+ """
+ prefix should be module_name to make keys identical to sequetial ones.
+ """
+ sharded_state_dict = {}
+ full_state_dict = self.state_dict(prefix='', keep_vars=True)
+ num_global_experts = get_expert_model_parallel_world_size() * self.num_gemms
+ local_expert_indices_offset = get_expert_model_parallel_rank() * self.num_gemms
+ ep_axis = len(sharded_offsets)
+ extra_states = self._split_extra_state(full_state_dict['_extra_state'])
+ for gemm_idx in range(self.num_gemms):
+ state_dict = {
+ f'{gemm_idx}.weight': full_state_dict[f'weight{gemm_idx}'],
+ f'{gemm_idx}._extra_state': extra_states[gemm_idx],
+ }
+ if self.use_bias:
+ state_dict[f'{gemm_idx}.bias'] = full_state_dict[f'bias{gemm_idx}']
+ sub_sd = make_sharded_tensors_for_checkpoint(
+ state_dict,
+ '',
+ tp_axis_map,
+ (
+ *sharded_offsets,
+ (ep_axis, local_expert_indices_offset + gemm_idx, num_global_experts),
+ ),
+ )
+ # Remove expert layers indexing from sharded keys
+ replace_prefix_for_sharding(sub_sd, f'{gemm_idx}.', prefix)
+ sharded_state_dict.update(
+ {
+ f'{prefix}weight{gemm_idx}': sub_sd[f'{gemm_idx}.weight'],
+ f'{prefix}_extra_state{"" if gemm_idx == 0 else gemm_idx}': sub_sd[
+ f'{gemm_idx}._extra_state'
+ ],
+ }
+ )
+ if self.use_bias:
+ sharded_state_dict[f'{prefix}bias{gemm_idx}'] = sub_sd[f'{gemm_idx}.bias']
+ # Adjust replica ids - replication along DP modulo EP
+ for k, sh_ten in sharded_state_dict.items():
+ replica_id = sh_ten.replica_id
+ assert (
+ len(replica_id) == 3
+ ), f'Expected replica_id for {k} to be in (PP, TP, DP) format, got: {replica_id}'
+ if getattr(sh_ten, "is_data_parallel_fully_shard", False):
+ edp_replica_id = 0
+ else:
+ edp_replica_id = get_expert_data_parallel_rank()
+ sh_ten.replica_id = (*replica_id[:2], edp_replica_id)
+ return sharded_state_dict
+
+ def backward_dw(self):
+ """Compute weight gradients during the backward pass if split_bw is enabled."""
+ if self.config.split_bw:
+ super().backward_dw()
+
+ class TEColumnParallelGroupedLinear(TEGroupedLinear):
+ """
+ Wrapper for the Transformer-Engine's `GroupedLinear` layer but specialized
+ to column-parallel style.
+ """
+
+ def __init__(
+ self,
+ num_gemms: int,
+ input_size: int,
+ output_size: int,
+ *,
+ config: ModelParallelConfig,
+ init_method: Callable,
+ bias: bool,
+ skip_bias_add: bool,
+ is_expert: bool,
+ tp_comm_buffer_name: Optional[str] = None,
+ ):
+
+ super().__init__(
+ num_gemms=num_gemms,
+ input_size=input_size,
+ output_size=output_size,
+ parallel_mode="column",
+ config=config,
+ init_method=condition_init_method(config, init_method),
+ bias=bias,
+ skip_bias_add=skip_bias_add,
+ is_expert=is_expert,
+ tp_comm_buffer_name=tp_comm_buffer_name,
+ )
+
+ def sharded_state_dict(self, prefix='', sharded_offsets=(), metadata=None):
+ """
+ For each gemm, sharding along axis 0, bias sharded.
+ Assume sharded_offsets[-1] is the expert parallel offset.
+ """
+ tp_axis_map = {}
+ for gemm_idx in range(self.num_gemms):
+ tp_axis_map.update({f'{gemm_idx}.weight': 0, f'{gemm_idx}.bias': 0})
+ return super()._sharded_state_dict_grouped(
+ tp_axis_map, prefix, sharded_offsets, metadata
+ )
+
+ class TERowParallelGroupedLinear(TEGroupedLinear):
+ """
+ Wrapper for the Transformer-Engine's `GroupedLinear` layer but specialized
+ to row-parallel style.
+ """
+
+ def __init__(
+ self,
+ num_gemms: int,
+ input_size: int,
+ output_size: int,
+ *,
+ config: ModelParallelConfig,
+ init_method: Callable,
+ bias: bool,
+ skip_bias_add: bool,
+ is_expert: bool,
+ tp_comm_buffer_name: Optional[str] = None,
+ ):
+
+ super().__init__(
+ num_gemms=num_gemms,
+ input_size=input_size,
+ output_size=output_size,
+ parallel_mode="row",
+ config=config,
+ init_method=condition_init_method(config, init_method),
+ bias=bias,
+ skip_bias_add=skip_bias_add,
+ is_expert=is_expert,
+ tp_comm_buffer_name=tp_comm_buffer_name,
+ )
+
+ def sharded_state_dict(self, prefix='', sharded_offsets=(), metadata=None):
+ """
+ For each gemm, sharding along axis 1, bias not sharded.
+ Assume sharded_offsets[-1] is the expert parallel offset.
+ """
+ tp_axis_map = {f'{gemm_idx}.weight': 1 for gemm_idx in range(self.num_gemms)}
+ return super()._sharded_state_dict_grouped(
+ tp_axis_map, prefix, sharded_offsets, metadata
+ )
+
+else:
+ TEGroupedLinear = None # type: ignore[assignment, misc]
+ TEColumnParallelGroupedLinear = None # type: ignore[assignment, misc]
+ TERowParallelGroupedLinear = None # type: ignore[assignment, misc]
+
+
+class TEDelayedScaling(te.common.recipe.DelayedScaling):
+ """
+ Wrapper for the Transformer-Engine's `DelayedScaling` layer.
+ """
+
+ def __init__(
+ self,
+ config: ModelParallelConfig,
+ fp8_format: int,
+ override_linear_precision: tuple = (False, False, False),
+ ):
+ extra_kwargs = _get_extra_te_kwargs(config)
+ if is_te_min_version("1.6.0.dev0"):
+ extra_kwargs["fp8_dpa"] = config.fp8_dot_product_attention
+ extra_kwargs["fp8_mha"] = config.fp8_multi_head_attention
+ if get_te_version() < PkgVersion("1.8.0"):
+ extra_kwargs["interval"] = config.fp8_interval
+ elif config.fp8_interval != 1:
+ warnings.warn("fp8_interval is deprecated and ignored from Transformer-Engine v1.8.0.")
+
+ super().__init__(
+ margin=config.fp8_margin,
+ fp8_format=fp8_format,
+ amax_compute_algo=config.fp8_amax_compute_algo,
+ amax_history_len=config.fp8_amax_history_len,
+ override_linear_precision=override_linear_precision,
+ **extra_kwargs,
+ )
+
+
+class TECudaRNGStatesTracker(te.pytorch.distributed.CudaRNGStatesTracker):
+ """Wraps TransformerEngine's CudaRNGStatesTracker so that it is
+ interchangeable with Megatron's RNG tracker"""
+
+ def __init__(self):
+ super().__init__()
+ self.reset()
+
+ def is_initialized(self):
+ """Checks if the internal RNG state has been set wirth set_states()."""
+ return self._is_initialized
+
+ def reset(self):
+ """Reset the internal RNG state."""
+ super().reset()
+ self._is_initialized = False
+
+ def set_states(self, states):
+ """Set the internal RNG state."""
+ super().set_states(states)
+ self._is_initialized = True
+
+ def add(self, name, seed):
+ """Track the rng state."""
+ super().add(name, seed)
+ self._is_initialized = True
+
+
+def te_checkpoint(
+ forward_func,
+ distribute_saved_activations,
+ get_rng_state_tracker,
+ tp_group,
+ hidden_states,
+ attention_mask,
+ attn_mask_type,
+ context,
+ context_mask,
+ rotary_pos_emb,
+ **kwargs,
+):
+ """Checkpointing with Transformer-Engine."""
+ from transformer_engine.pytorch.distributed import checkpoint
+
+ if is_te_min_version("1.5.0"):
+ return checkpoint(
+ forward_func,
+ hidden_states,
+ attention_mask,
+ attn_mask_type,
+ context,
+ context_mask,
+ rotary_pos_emb,
+ distribute_saved_activations=distribute_saved_activations,
+ get_rng_state_tracker=get_rng_state_tracker,
+ tp_group=tp_group,
+ **kwargs,
+ )
+ else:
+ return checkpoint(
+ forward_func,
+ distribute_saved_activations,
+ get_rng_state_tracker,
+ tp_group,
+ hidden_states,
+ attention_mask,
+ context,
+ context_mask,
+ rotary_pos_emb,
+ )
+
+
+try:
+ from transformer_engine.pytorch.attention import _SplitAlongDim
+ SplitAlongDim = _SplitAlongDim.apply
+
+except ImportError:
+ SplitAlongDim = None
+
+try:
+ from transformer_engine.pytorch.cpu_offload import get_cpu_offload_context as _get_cpu_offload_context
+ def get_cpu_offload_context(
+ enabled, num_layers, model_layers, activation_offloading, weight_offloading
+ ):
+ """Get CPU offload context and sync function."""
+ if is_te_min_version("1.10.0.dev0"):
+ context, sync_func = _get_cpu_offload_context(
+ enabled, num_layers, model_layers, activation_offloading, weight_offloading
+ )
+ else:
+ context, sync_func = _get_cpu_offload_context(
+ enabled, num_layers, activation_offloading, weight_offloading
+ )
+
+ return context, sync_func
+
+except ImportError:
+ get_cpu_offload_context = None # type: ignore[assignment, misc]
+
+try:
+ from transformer_engine.pytorch.attention import FusedRoPEFunc
+ def fused_apply_rotary_pos_emb(
+ t: torch.Tensor, freqs: torch.Tensor, transpose_output_memory: bool = False
+ ) -> torch.Tensor:
+ """Apply rotary positional embedding to input tensor T in `sbhd` format."""
+ if transpose_output_memory:
+ warnings.warn(
+ "transpose_output_memory is not supported by TE's fused RoPE and will be ignored."
+ )
+ return FusedRoPEFunc.apply(t, freqs, "sbhd")
+
+ def fused_apply_rotary_pos_emb_thd(
+ t: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ freqs: torch.Tensor,
+ cp_size: int = 1,
+ cp_rank: int = 0,
+ ) -> torch.Tensor:
+ """
+ Apply rotary positional embedding to input tensor T in `thd` format with CP support.
+ """
+ if is_te_min_version("1.11.0", check_equality=False):
+ return FusedRoPEFunc.apply(t, freqs, "thd", cu_seqlens, cp_size, cp_rank)
+ else:
+ return FusedRoPEFunc.apply(t, freqs, "thd", cu_seqlens)
+
+except ImportError:
+ pass
+
+try:
+ from transformer_engine.pytorch import Fp8Padding, Fp8Unpadding # pylint: disable=unused-import
+except ImportError:
+ Fp8Padding = None
+ Fp8Unpadding = None
+
+try:
+ from transformer_engine.pytorch.permutation import (
+ moe_permute,
+ moe_sort_chunks_by_index,
+ moe_unpermute,
+ )
+ fused_permute = moe_permute
+ fused_unpermute = moe_unpermute
+ fused_sort_chunks_by_index = moe_sort_chunks_by_index
+
+except ImportError:
+ fused_permute = None
+ fused_unpermute = None
+ fused_sort_chunks_by_index = None
\ No newline at end of file
diff --git a/aiak_megatron/megatron/core/fp8_utils.py b/aiak_megatron/megatron/core/fp8_utils.py
index 76db15b0..1cd3814a 100644
--- a/aiak_megatron/megatron/core/fp8_utils.py
+++ b/aiak_megatron/megatron/core/fp8_utils.py
@@ -323,10 +323,14 @@ def _correct_amax_history_if_needed_impl(model: List[torch.nn.Module]) -> None:
else:
# Fallback impl if TE version is invalid or TE is not installed.
def _modify_underlying_storage_impl(*args, **kwargs):
- raise RuntimeError("Invalid Transformer Engine version for FP8 distributed optimizer")
+ # raise RuntimeError("Invalid Transformer Engine version for FP8 distributed optimizer")
+ # No-op fallback: if TE is not installed, we're not using FP8, so just skip
+ pass
def _quantize_param_shard_impl(*args, **kwargs):
- raise RuntimeError("Invalid Transformer Engine version for FP8 distributed optimizer")
+ # raise RuntimeError("Invalid Transformer Engine version for FP8 distributed optimizer")
+ # No-op fallback: if TE is not installed, we're not using FP8, so just skip
+ pass
def _correct_amax_history_if_needed_impl(*args, **kwargs):
# If TE is not installed, we are definitely not using fp8 for training, so no correction
diff --git a/aiak_megatron/megatron/core/fusions/fused_layer_norm.py b/aiak_megatron/megatron/core/fusions/fused_layer_norm.py
index 01a654db..0bcd32b6 100644
--- a/aiak_megatron/megatron/core/fusions/fused_layer_norm.py
+++ b/aiak_megatron/megatron/core/fusions/fused_layer_norm.py
@@ -8,25 +8,25 @@
from torch import Tensor
from torch.nn import init
from torch.nn.parameter import Parameter
-
+import torch.nn.functional as F
from megatron.core.transformer import TransformerConfig
from megatron.core.utils import make_viewless_tensor
-try:
- from apex.contrib.layer_norm.layer_norm import FastLayerNormFN
-
- HAVE_PERSIST_LAYER_NORM = True
-except ImportError:
- HAVE_PERSIST_LAYER_NORM = False
+# try:
+# from apex.contrib.layer_norm.layer_norm import FastLayerNormFN
-try:
- from apex.normalization.fused_layer_norm import FusedLayerNormAffineFunction
-
- HAVE_FUSED_LAYER_NORM = True
-except ImportError:
- HAVE_FUSED_LAYER_NORM = False
+# HAVE_PERSIST_LAYER_NORM = True
+# except ImportError:
+# HAVE_PERSIST_LAYER_NORM = False
+# try:
+# from apex.normalization.fused_layer_norm import FusedLayerNormAffineFunction
+# HAVE_FUSED_LAYER_NORM = True
+# except ImportError:
+# HAVE_FUSED_LAYER_NORM = False
+HAVE_PERSIST_LAYER_NORM = False
+HAVE_FUSED_LAYER_NORM = False
class FusedLayerNorm(torch.nn.Module):
"""Layer Norm, fused into a single CUDA kernel.
@@ -102,7 +102,12 @@ def __init__(
if not persist_layer_norm and not HAVE_FUSED_LAYER_NORM:
# TODO: Add pytorch only layer norm
- raise ValueError(f'Apex must be installed to use FusedLayerNorm.')
+ self._use_pytorch_fallback = True
+ persist_layer_norm = False
+ else:
+ self._use_pytorch_fallback = False
+
+
if isinstance(hidden_size, numbers.Integral):
hidden_size = (hidden_size,)
@@ -130,8 +135,15 @@ def reset_parameters(self):
init.zeros_(self.bias)
def forward(self, input: Tensor) -> Tensor:
-
+
weight = self.weight + 1 if self.zero_centered_gamma else self.weight
+ # If we flagged the pure-PyTorch fallback, use torch.nn.functional.layer_norm
+ if getattr(self, "_use_pytorch_fallback", False):
+ # weight should be offset when zero_centered_gamma is used (matches original behavior)
+ return F.layer_norm(input, tuple(self.hidden_size), weight=weight, bias=self.bias, eps=self.eps)
+
+
+
if self.persist_layer_norm:
if 'memory_efficient' in inspect.getfullargspec(FastLayerNormFN.forward).args:
diff --git a/aiak_megatron/megatron/core/fusions/fused_layer_norm_org.py b/aiak_megatron/megatron/core/fusions/fused_layer_norm_org.py
new file mode 100644
index 00000000..f6d472d6
--- /dev/null
+++ b/aiak_megatron/megatron/core/fusions/fused_layer_norm_org.py
@@ -0,0 +1,170 @@
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+
+import importlib
+import inspect
+import numbers
+
+import torch
+from torch import Tensor
+from torch.nn import init
+from torch.nn.parameter import Parameter
+
+from megatron.core.transformer import TransformerConfig
+from megatron.core.utils import make_viewless_tensor
+
+# try:
+# from apex.contrib.layer_norm.layer_norm import FastLayerNormFN
+
+# HAVE_PERSIST_LAYER_NORM = True
+# except ImportError:
+# HAVE_PERSIST_LAYER_NORM = False
+
+# try:
+# from apex.normalization.fused_layer_norm import FusedLayerNormAffineFunction
+
+# HAVE_FUSED_LAYER_NORM = True
+# except ImportError:
+# HAVE_FUSED_LAYER_NORM = False
+HAVE_PERSIST_LAYER_NORM = False
+HAVE_FUSED_LAYER_NORM = False
+class FusedLayerNorm(torch.nn.Module):
+ """Layer Norm, fused into a single CUDA kernel.
+
+ Args:
+ hidden_size (int): Transformer hidden dimension.
+
+ eps (float): Epsilon added to denominator, for numerical stability.
+
+ persist_layer_norm (bool): Use persistent fused layer norm kernel.
+ This kernel supports only a set of hidden sizes. Please
+ check persist_ln_hidden_sizes if your hidden size is supported.
+
+ zero_centered_gamma (bool): Adjust LayerNorm weights such that they are
+ centered around zero. This improves numerical stability.
+
+ config (TransformerConfig): Transformer config. Include to match custom
+ layer norm interfaces.
+
+ normalization (str): Normalization type, used for Transformer Engine.
+ Must equal 'LayerNorm' here.
+ """
+
+ def __init__(
+ self,
+ config: TransformerConfig,
+ hidden_size: int,
+ eps: float = 1e-5,
+ persist_layer_norm: bool = True,
+ zero_centered_gamma: bool = False,
+ normalization: str = "LayerNorm", # included to match TE interface
+ ):
+ super().__init__()
+
+ self.config = config
+
+ self.zero_centered_gamma = self.config.layernorm_zero_centered_gamma
+ assert (
+ self.config.normalization == "LayerNorm"
+ ), f'({self.config.normalization}) is not supported in FusedLayerNorm'
+
+ # List of hiddens sizes supported in the persistent layer norm kernel
+ # If the hidden size is not supported, fall back to the non-persistent
+ # kernel.
+ persist_ln_hidden_sizes = [
+ 1024,
+ 1536,
+ 2048,
+ 2304,
+ 3072,
+ 3840,
+ 4096,
+ 5120,
+ 6144,
+ 8192,
+ 10240,
+ 12288,
+ 12800,
+ 15360,
+ 16384,
+ 18432,
+ 20480,
+ 24576,
+ 25600,
+ 30720,
+ 32768,
+ 40960,
+ 49152,
+ 65536,
+ ]
+ persist_layer_norm = self.config.persist_layer_norm
+ if hidden_size not in persist_ln_hidden_sizes or not HAVE_PERSIST_LAYER_NORM:
+ persist_layer_norm = False
+
+ if not persist_layer_norm and not HAVE_FUSED_LAYER_NORM:
+ # TODO: Add pytorch only layer norm
+ raise ValueError(f'Apex must be installed to use FusedLayerNorm.')
+
+ if isinstance(hidden_size, numbers.Integral):
+ hidden_size = (hidden_size,)
+ self.hidden_size = torch.Size(hidden_size)
+ self.eps = eps
+ # Parameters need to be initialized with torch.empty rather than torch.Tensor
+ # for correct device placement with nemo2.
+ self.weight = Parameter(torch.empty(*hidden_size))
+ self.bias = Parameter(torch.empty(*hidden_size))
+ self.reset_parameters()
+ self.persist_layer_norm = persist_layer_norm
+ self.sequence_parallel = self.config.sequence_parallel
+
+ # set sequence parallelism flag on weight and bias parameters
+ setattr(self.weight, 'sequence_parallel', self.sequence_parallel)
+ setattr(self.bias, 'sequence_parallel', self.sequence_parallel)
+
+ def reset_parameters(self):
+
+ if self.zero_centered_gamma:
+ init.zeros_(self.weight)
+ init.zeros_(self.bias)
+ else:
+ init.ones_(self.weight)
+ init.zeros_(self.bias)
+
+ def forward(self, input: Tensor) -> Tensor:
+
+ weight = self.weight + 1 if self.zero_centered_gamma else self.weight
+
+ if self.persist_layer_norm:
+ if 'memory_efficient' in inspect.getfullargspec(FastLayerNormFN.forward).args:
+ output = FastLayerNormFN.apply(
+ input, weight, self.bias, self.eps, self.config.memory_efficient_layer_norm
+ )
+ else:
+ output = FastLayerNormFN.apply(input, weight, self.bias, self.eps)
+
+ # Apex's fast layer norm function outputs a 'view' tensor (i.e., has
+ # a populated '_base' field). This will result in schedule.py's
+ # deallocate_output_tensor() throwing an error, so a viewless tensor is
+ # created to prevent this.
+ output = make_viewless_tensor(
+ inp=output, requires_grad=input.requires_grad, keep_graph=True
+ )
+
+ else:
+ if (
+ 'memory_efficient'
+ in inspect.getfullargspec(FusedLayerNormAffineFunction.forward).args
+ ):
+ return FusedLayerNormAffineFunction.apply(
+ input,
+ weight,
+ self.bias,
+ self.hidden_size,
+ self.eps,
+ self.config.memory_efficient_layer_norm,
+ )
+ else:
+ return FusedLayerNormAffineFunction.apply(
+ input, weight, self.bias, self.hidden_size, self.eps
+ )
+
+ return output
diff --git a/aiak_megatron/megatron/core/model_parallel_config.py b/aiak_megatron/megatron/core/model_parallel_config.py
index a0022a08..a89fc3d6 100644
--- a/aiak_megatron/megatron/core/model_parallel_config.py
+++ b/aiak_megatron/megatron/core/model_parallel_config.py
@@ -143,6 +143,7 @@ class ModelParallelConfig:
# Optimizations
###################
gradient_accumulation_fusion: bool = False
+ print(gradient_accumulation_fusion)
"""If true, fuses weight gradient accumulation to GEMMs. Requires the custom CUDA extension
fused_weight_gradient_mlp_cuda module. To use gradient_accumulation_fusion you must install
APEX with --cpp_ext and --cuda_ext. For example: "pip install --global-option=\"--cpp_ext\"
diff --git a/aiak_megatron/megatron/core/models/common/embeddings/rope_utils.py b/aiak_megatron/megatron/core/models/common/embeddings/rope_utils.py
index 6c37d48a..7d0de0fa 100644
--- a/aiak_megatron/megatron/core/models/common/embeddings/rope_utils.py
+++ b/aiak_megatron/megatron/core/models/common/embeddings/rope_utils.py
@@ -186,6 +186,12 @@ def apply_rotary_pos_emb(
"""
Reroute to the appropriate apply_rotary_pos_emb function depending on
fused/unfused kernels, or bshd (conventional) / thd (packed seq) format
+ # DEBUG: Check input
+ if t is None:
+ print(f"[apply_rotary_pos_emb] ERROR: input t is None!", flush=True)
+ return None
+
+ print(f"[apply_rotary_pos_emb] config.apply_rope_fusion={config.apply_rope_fusion}, cu_seqlens={cu_seqlens is not None}, fused available={fused_apply_rotary_pos_emb is not None}", flush=True)
"""
if config.apply_rope_fusion and config.multi_latent_attention is False:
diff --git a/aiak_megatron/megatron/core/optimizer/__init__.py b/aiak_megatron/megatron/core/optimizer/__init__.py
index e4a38cbd..d52bd68e 100644
--- a/aiak_megatron/megatron/core/optimizer/__init__.py
+++ b/aiak_megatron/megatron/core/optimizer/__init__.py
@@ -6,22 +6,23 @@
import torch
from torch.optim import SGD as CPUSGD
from torch.optim import AdamW as CPUAdam
-
-try:
- from transformer_engine.pytorch.optimizers import FusedAdam as Adam
- from transformer_engine.pytorch.optimizers import FusedSGD as SGD
-except ImportError:
- try:
- from apex.optimizers import FusedAdam as Adam
- from apex.optimizers import FusedSGD as SGD
- except ImportError:
- warnings.warn(f'Transformer Engine and Apex are not installed. Falling back to Torch optimizers.')
-
- # Apex's FusedAdam is a drop-in replacement for torch's AdamW.
- # pylint: disable-next=line-too-long.
- # See https://github.com/NVIDIA/apex/blob/7b73b12361068a10b0f44844534613f252a5ea75/apex/optimizers/fused_adam.py#L16.
- from torch.optim import AdamW as Adam, SGD
-
+# from megatron.core.optimizer import Adam as Adam
+# try:
+# from transformer_engine.pytorch.optimizers import FusedAdam as Adam
+# from transformer_engine.pytorch.optimizers import FusedSGD as SGD
+# except ImportError:
+# try:
+# from apex.optimizers import FusedAdam as Adam
+# from apex.optimizers import FusedSGD as SGD
+# except ImportError:
+# warnings.warn(f'Transformer Engine and Apex are not installed. Falling back to Torch optimizers.')
+
+# # Apex's FusedAdam is a drop-in replacement for torch's AdamW.
+# # pylint: disable-next=line-too-long.
+# # See https://github.com/NVIDIA/apex/blob/7b73b12361068a10b0f44844534613f252a5ea75/apex/optimizers/fused_adam.py#L16.
+# from torch.optim import AdamW as Adam, SGD
+Adam = CPUAdam
+SGD = CPUSGD
from megatron.core import mpu
from megatron.core.optimizer.cpu_offloading.hybrid_optimizer import HybridDeviceOptimizer
diff --git a/aiak_megatron/megatron/core/optimizer/distrib_optimizer.py b/aiak_megatron/megatron/core/optimizer/distrib_optimizer.py
index 6b565cee..a006fb35 100644
--- a/aiak_megatron/megatron/core/optimizer/distrib_optimizer.py
+++ b/aiak_megatron/megatron/core/optimizer/distrib_optimizer.py
@@ -10,16 +10,19 @@
import torch
-HAVE_APEX_OR_TE = True
-try:
- from transformer_engine.pytorch.optimizers import FusedAdam as Adam
-except ImportError:
- try:
- from apex.optimizers import FusedAdam as Adam
- except ImportError:
- from torch.optim import AdamW as Adam
+# HAVE_APEX_OR_TE = True
+# try:
+# from transformer_engine.pytorch.optimizers import FusedAdam as Adam
+# except ImportError:
+# try:
+# from apex.optimizers import FusedAdam as Adam
+# except ImportError:
+# from torch.optim import AdamW as Adam
- HAVE_APEX_OR_TE = False
+# HAVE_APEX_OR_TE = False
+from torch.optim import AdamW as Adam
+
+HAVE_APEX_OR_TE = False
from megatron.core.optimizer.cpu_offloading import HybridDeviceOptimizer
@@ -486,7 +489,7 @@ def __init__(
for model_chunk in self.model_chunks:
assert self.ddp_config == model_chunk.ddp_config
self.distributed_optimizer_instance_id = distributed_optimizer_instance_id
-
+ print(optimizer)
assert isinstance(optimizer, (Adam, HybridDeviceOptimizer)) or optimizer is None, (
"Only Adam and HybridDeviceOptimizer currently supported, due to checkpointing requirements."
)
diff --git a/aiak_megatron/megatron/core/tensor_parallel/layers.py b/aiak_megatron/megatron/core/tensor_parallel/layers.py
index 8689602a..ff1b8d67 100644
--- a/aiak_megatron/megatron/core/tensor_parallel/layers.py
+++ b/aiak_megatron/megatron/core/tensor_parallel/layers.py
@@ -854,7 +854,7 @@ def __init__(
"gradient accumulation fusion."
)
self.gradient_accumulation_fusion = config.gradient_accumulation_fusion
-
+
if self.allreduce_dgrad and self.sequence_parallel:
raise RuntimeError(
"`allreduce_dgrad` and `sequence_parallel` cannot be enabled at the same time."
@@ -938,7 +938,7 @@ def forward(
self._forward_impl = linear_with_grad_accumulation_and_async_allreduce
allreduce_dgrad = False if self.explicit_expert_comm else self.allreduce_dgrad
-
+
output_parallel = self._forward_impl(
input=input_parallel,
weight=weight,
diff --git a/aiak_megatron/megatron/core/transformer/attention.py b/aiak_megatron/megatron/core/transformer/attention.py
index b489a35b..c034fd59 100644
--- a/aiak_megatron/megatron/core/transformer/attention.py
+++ b/aiak_megatron/megatron/core/transformer/attention.py
@@ -106,6 +106,16 @@ def __init__(
else:
self.apply_rotary_fn = None
+ # Log which rotary function is bound (once for early layer to avoid spam)
+ if self.apply_rotary_fn is not None and getattr(self, 'layer_number', 0) in [0, 1]:
+ try:
+ fn = self.apply_rotary_fn
+ fn_name = getattr(fn, '__name__', type(fn).__name__)
+ fn_mod = getattr(fn, '__module__', str(fn))
+ print(f"[Attention] apply_rotary_fn bound to: {fn_mod}.{fn_name}", flush=True)
+ except Exception as e:
+ print(f"[Attention] apply_rotary_fn identification failed: {e}", flush=True)
+
# To support both CUDA Graphs and key value with different hidden size
self.key_hidden_size = self.hidden_size_per_attention_head
self.val_hidden_size = self.hidden_size_per_attention_head
@@ -378,7 +388,7 @@ def forward(
# self or cross attn.
query, key, value = self.get_query_key_value_tensors(hidden_states, key_value_states)
-
+
# ===================================================
# Adjust key, value, and rotary_pos_emb for inference
# ===================================================
@@ -421,6 +431,10 @@ def forward(
attn_mask_type,
sequence_len_offset,
)
+
+ # DEBUG: Check query/key/value after _adjust_key_value_for_inference
+ if query is None or key is None:
+ print(f"[Attention] WARNING: query/key became None after _adjust_key_value_for_inference! query={query is not None}, key={key is not None}, value={value is not None}", flush=True)
if packed_seq_params is not None:
query = query.squeeze(1)
@@ -446,8 +460,30 @@ def forward(
cu_seqlens_q = cu_seqlens_kv = None
assert self.apply_rotary_fn is not None, "apply_rotary_fn must be defined"
- query = self.apply_rotary_fn(query, q_pos_emb, config=self.config, cu_seqlens=cu_seqlens_q)
- key = self.apply_rotary_fn(key, k_pos_emb, config=self.config, cu_seqlens=cu_seqlens_kv)
+ # Keep originals as fallback in case RoPE function misbehaves
+ _orig_query = query
+ _orig_key = key
+ try:
+ query = self.apply_rotary_fn(query, q_pos_emb, config=self.config, cu_seqlens=cu_seqlens_q)
+ if query is None:
+ print(f"[Attention] ERROR: apply_rotary_fn returned None for query! Using unrotated query.", flush=True)
+ query = _orig_query
+ except Exception as e:
+ print(f"[Attention] EXCEPTION in apply_rotary_fn(query): {e}. Using unrotated query.", flush=True)
+ query = _orig_query
+
+ try:
+ key = self.apply_rotary_fn(key, k_pos_emb, config=self.config, cu_seqlens=cu_seqlens_kv)
+ if key is None:
+ print(f"[Attention] ERROR: apply_rotary_fn returned None for key! Using unrotated key.", flush=True)
+ key = _orig_key
+ except Exception as e:
+ print(f"[Attention] EXCEPTION in apply_rotary_fn(key): {e}. Using unrotated key.", flush=True)
+ key = _orig_key
+
+ # DEBUG: Check query/key after apply_rotary_fn
+ if query is None or key is None:
+ print(f"[Attention] WARNING: query/key became None after apply_rotary_fn! query={query is not None}, key={key is not None}", flush=True)
# TODO, can apply positional embedding to value_layer so it has
# absolute positional embedding.
@@ -457,6 +493,14 @@ def forward(
# ==================================
# core attention computation
# ==================================
+
+ # Safety check: if query, key, or value are None, return a zero tensor with matching shape
+ if query is None or key is None or value is None:
+ print(f"[Attention] WARNING: query/key/value is None before core_attention! query={query is not None}, key={key is not None}, value={value is not None}", flush=True)
+ # Return zeros with expected shape matching hidden_states [sq, b, hidden_size]
+ # hidden_states shape should be preserved through the projection
+ return (torch.zeros_like(hidden_states), None)
+
if self.checkpoint_core_attention and self.training:
core_attn_out = self._checkpointed_attention_forward(
query,
@@ -629,26 +673,39 @@ def get_query_key_value_tensors(self, hidden_states, key_value_states=None):
# Attention heads [sq, b, h] --> [sq, b, ng * (np/ng + 2) * hn)]
mixed_qkv, _ = self.linear_qkv(hidden_states)
- # [sq, b, hp] --> [sq, b, ng, (np/ng + 2) * hn]
- new_tensor_shape = mixed_qkv.size()[:-1] + (
- self.num_query_groups_per_partition,
- (
- (self.num_attention_heads_per_partition // self.num_query_groups_per_partition + 2)
- * self.hidden_size_per_attention_head
- ),
- )
- mixed_qkv = mixed_qkv.view(*new_tensor_shape)
-
- split_arg_list = [
- (
- self.num_attention_heads_per_partition
- // self.num_query_groups_per_partition
- * self.hidden_size_per_attention_head
- ),
- self.hidden_size_per_attention_head,
- self.hidden_size_per_attention_head,
- ]
-
+ # [sq, b, hp] --> [sq, b, ng, inferred_group_width]
+ # old static shape assumption kept for reference:
+ # new_tensor_shape = mixed_qkv.size()[:-1] + (
+ # self.num_query_groups_per_partition,
+ # (
+ # (self.num_attention_heads_per_partition // self.num_query_groups_per_partition + 2)
+ # * self.hidden_size_per_attention_head
+ # ),
+ # )
+ # mixed_qkv = mixed_qkv.view(*new_tensor_shape)
+
+ ng = self.num_query_groups_per_partition
+ last_dim = mixed_qkv.size(-1)
+ group_width = last_dim // ng
+ mixed_qkv = mixed_qkv.reshape(mixed_qkv.size(0), mixed_qkv.size(1), ng, group_width)
+
+ q_width = (
+ self.num_attention_heads_per_partition // self.num_query_groups_per_partition
+ ) * self.hidden_size_per_attention_head
+ remaining = group_width - q_width
+ # if the original assumption does not hold, split the remainder evenly for k and v
+ kv_width = max(0, remaining // 2)
+ split_arg_list = [q_width, kv_width, group_width - q_width - kv_width]
+
+ # DEBUG: Check if any split size is 0
+ if any(s == 0 for s in split_arg_list):
+ print(f"[get_query_key_value_tensors] WARNING: Zero-sized split detected!", flush=True)
+ print(f" ng={ng}, group_width={group_width}, q_width={q_width}, kv_width={kv_width}", flush=True)
+ print(f" split_arg_list={split_arg_list}", flush=True)
+ print(f" num_attention_heads_per_partition={self.num_attention_heads_per_partition}", flush=True)
+ print(f" num_query_groups_per_partition={self.num_query_groups_per_partition}", flush=True)
+ print(f" hidden_size_per_attention_head={self.hidden_size_per_attention_head}", flush=True)
+
if SplitAlongDim is not None:
# [sq, b, ng, (np/ng + 2) * hn]
# --> [sq, b, ng, np/ng * hn], [sq, b, ng, hn], [sq, b, ng, hn]
@@ -658,6 +715,14 @@ def get_query_key_value_tensors(self, hidden_states, key_value_states=None):
# --> [sq, b, ng, np/ng * hn], [sq, b, ng, hn], [sq, b, ng, hn]
(query, key, value) = torch.split(mixed_qkv, split_arg_list, dim=3)
+ # DEBUG: Check if split produced None or empty tensors
+ if query is None or query.numel() == 0:
+ print(f"[get_query_key_value_tensors] ERROR: query is None or empty! query={query}", flush=True)
+ if key is None or key.numel() == 0:
+ print(f"[get_query_key_value_tensors] ERROR: key is None or empty! key={key}", flush=True)
+ if value is None or value.numel() == 0:
+ print(f"[get_query_key_value_tensors] ERROR: value is None or empty! value={value}", flush=True)
+
# [sq, b, ng, np/ng * hn] -> [sq, b, np, hn]
query = query.reshape(query.size(0), query.size(1), -1, self.hidden_size_per_attention_head)
diff --git a/aiak_megatron/megatron/core/transformer/dot_product_attention.py b/aiak_megatron/megatron/core/transformer/dot_product_attention.py
index 1b5200ae..dcd79346 100644
--- a/aiak_megatron/megatron/core/transformer/dot_product_attention.py
+++ b/aiak_megatron/megatron/core/transformer/dot_product_attention.py
@@ -16,6 +16,71 @@
from megatron.core.transformer.utils import attention_mask_func
from megatron.core.utils import divide
+import torch.nn.functional as F
+from typing import Optional
+
+class ScaleMaskSoftmaxFallback:
+ """
+ Pure-PyTorch fallback for fused scale+mask+softmax.
+
+ Parameters:
+ - softmax_in_fp32: if True, compute softmax in float32 for stability.
+ - attn_mask_type: (optional) unused here except for API compatibility.
+ - scale: a scalar or None; FusedScaleMaskSoftmax used 'scale' for layer scaling
+ in some implementations. This fallback will NOT multiply or divide
+ by `scale` automatically (because your code already applies `self.softmax_scale`
+ in the matmul). If you need special handling, you can pass scale and
+ we will multiply scores by scale here.
+ """
+
+ def __init__(self, softmax_in_fp32: bool = False, attn_mask_type=None, scale: Optional[float] = None):
+ self.softmax_in_fp32 = softmax_in_fp32
+ self.attn_mask_type = attn_mask_type
+ self.scale = scale
+
+ def __call__(self, attention_scores: torch.Tensor, attention_mask: Optional[torch.Tensor], attn_mask_type: Optional[object] = None):
+ """
+ attention_scores: [b, np, sq, sk] (or equivalent)
+ attention_mask: broadcastable mask, or None. Convention in Megatron often
+ uses additive masks (large negative for masked positions).
+ Returns attention_probs: same shape as attention_scores
+ """
+ # Optionally apply an additional scale (some fused ops expect to apply layer scaling here).
+ if self.scale is not None:
+ # If scale is an integer layer_number in original fused implementation,
+ # you may want to divide or multiply depending on how self.softmax_scale was set.
+ # The original code sets self.softmax_scale = 1/sqrt(head_dim) and divides it by coeff
+ # (layer_number) if apply_query_key_layer_scaling is True. Since the matmul already
+ # used self.softmax_scale, we leave this as a no-op by default.
+ try:
+ # if user wants to apply the numeric scale here, do so:
+ attention_scores = attention_scores * float(self.scale)
+ except Exception:
+ pass
+
+ if attention_mask is not None:
+ # If attention_mask is additive (contains large negative values in masked positions),
+ # just add it. If it is boolean, convert it to additive mask.
+ if attention_mask.dtype == torch.bool:
+ # True = keep, False = masked? Depends on user's mask conventions.
+ # Typical convention: True indicates valid tokens -> we want mask where False positions are -inf.
+ # So convert boolean mask to additive: 0 for valid, -1e9 for masked.
+ add_mask = (~attention_mask).to(attention_scores.dtype) * -1e9
+ # Ensure add_mask broadcastable to attention_scores shape
+ attention_scores = attention_scores + add_mask
+ else:
+ # assume it's additive mask already shaped/broadcastable
+ attention_scores = attention_scores + attention_mask
+
+ # Softmax
+ if self.softmax_in_fp32:
+ orig_dtype = attention_scores.dtype
+ probs = torch.softmax(attention_scores.float(), dim=-1).to(orig_dtype)
+ else:
+ probs = torch.softmax(attention_scores, dim=-1)
+
+ return probs
+
class DotProductAttention(MegatronModule):
"""
@@ -78,16 +143,33 @@ def __init__(
coeff = self.layer_number
self.softmax_scale /= coeff
- self.scale_mask_softmax = FusedScaleMaskSoftmax(
- input_in_fp16=self.config.fp16,
- input_in_bf16=self.config.bf16,
- attn_mask_type=self.attn_mask_type,
- scaled_masked_softmax_fusion=self.config.masked_softmax_fusion,
- mask_func=attention_mask_func,
- softmax_in_fp32=self.config.attention_softmax_in_fp32,
- scale=coeff,
- )
-
+ # self.scale_mask_softmax = FusedScaleMaskSoftmax(
+ # input_in_fp16=self.config.fp16,
+ # input_in_bf16=self.config.bf16,
+ # attn_mask_type=self.attn_mask_type,
+ # scaled_masked_softmax_fusion=self.config.masked_softmax_fusion,
+ # mask_func=attention_mask_func,
+ # softmax_in_fp32=self.config.attention_softmax_in_fp32,
+ # scale=coeff,
+ # )
+ use_fused=False
+ if use_fused:
+ self.scale_mask_softmax = FusedScaleMaskSoftmax(
+ input_in_fp16=self.config.fp16,
+ input_in_bf16=self.config.bf16,
+ attn_mask_type=self.attn_mask_type,
+ scaled_masked_softmax_fusion=self.config.masked_softmax_fusion,
+ mask_func=attention_mask_func,
+ softmax_in_fp32=self.config.attention_softmax_in_fp32,
+ scale=coeff,
+ )
+ else:
+ # fallback implementation defined below (or import it)
+ self.scale_mask_softmax = ScaleMaskSoftmaxFallback(
+ softmax_in_fp32=self.config.attention_softmax_in_fp32,
+ attn_mask_type=self.attn_mask_type,
+ scale=coeff,
+ )
# Dropout. Note that for a single iteration, this layer will generate
# different outputs on different number of parallel partitions but
# on average it should not be partition dependent.
@@ -105,6 +187,7 @@ def forward(
attention_bias: Tensor = None,
packed_seq_params: Optional[PackedSeqParams] = None,
):
+ packed_seq_params=None
"""Forward."""
assert packed_seq_params is None, (
"Packed sequence is not supported by DotProductAttention."
@@ -123,23 +206,107 @@ def forward(
# attn_mask_type is not used.
if self.num_attention_heads_per_partition // self.num_query_groups_per_partition > 1:
- key = key.repeat_interleave(
- self.num_attention_heads_per_partition // self.num_query_groups_per_partition, dim=2
- )
- value = value.repeat_interleave(
- self.num_attention_heads_per_partition // self.num_query_groups_per_partition, dim=2
+ # Only repeat when tensors are 4D [*, b, ng, hn] so we expand ng -> np.
+ repeat_times = self.num_attention_heads_per_partition // self.num_query_groups_per_partition
+ if key is not None and key.dim() == 4:
+ key = key.repeat_interleave(repeat_times, dim=2)
+ if value is not None and value.dim() == 4:
+ value = value.repeat_interleave(repeat_times, dim=2)
+
+ # Normalize shapes to 3D for baddbmm and derive (b, np, sq, sk)
+ if query.dim() == 4:
+ # [sq, b, np, hn]
+ sq_len, b, np_per_part, hn = query.size()
+ sk_len = key.size(0)
+ # [sq, b, np, hn] -> [sq, b*np, hn]
+ query = query.reshape(sq_len, b * np_per_part, hn)
+ # Key can be 4D or 3D depending on upstream; normalize to 3D [sk, b*np, hn]
+ if key.dim() == 4:
+ key = key.reshape(sk_len, b * np_per_part, -1)
+ elif query.dim() == 3:
+ # [sq, b*np, hn]
+ sq_len, bnp, hn = query.size()
+ sk_len = key.size(0)
+ np_per_part = self.num_attention_heads_per_partition
+ assert bnp % np_per_part == 0, (
+ f"DotProductAttention: b*np ({bnp}) not divisible by heads_per_part ({np_per_part})"
)
+ b = bnp // np_per_part
+ else:
+ raise RuntimeError(f"DotProductAttention: Unsupported query.dim()={query.dim()}")
+
+ # Ensure key is 3D [sk, b*np, hn]
+ if key.dim() == 4:
+ # [sk, b, np, hn] -> [sk, b*np, hn]
+ key = key.reshape(sk_len, b * np_per_part, -1)
+ elif key.dim() == 3:
+ # If key is [sk, b*g, hn*r] for GQA packing, unflatten to [sk, b*np, hn]
+ if key.size(1) != b * np_per_part:
+ g = self.num_query_groups_per_partition
+ if key.size(1) == b * g:
+ r = key.size(2) // hn
+ if (r * g) == np_per_part and (key.size(2) % hn) == 0:
+ # [sk, b, g, r, hn] -> [sk, b, np, hn] -> [sk, b*np, hn]
+ key = (
+ key.reshape(sk_len, b, g, r, hn)
+ .permute(0, 1, 3, 2, 4)
+ .reshape(sk_len, b * np_per_part, hn)
+ )
+ elif np_per_part % g == 0 and r == 1:
+ # Simple case: [sk, b*g, hn] -> repeat groups to heads
+ repeat_times = np_per_part // g
+ key = key.repeat_interleave(repeat_times, dim=1)
+ else:
+ raise RuntimeError(
+ f"DotProductAttention: Cannot map key of shape {tuple(key.size())} to [sk, b*np, hn]. "
+ f"b={b}, np={np_per_part}, g={g}, hn={hn}"
+ )
+ elif key.size(1) == b:
+ key = key.repeat_interleave(np_per_part, dim=1)
+ else:
+ raise RuntimeError(
+ f"DotProductAttention: Unexpected key.size(1)={key.size(1)}; "
+ f"expected {b*np_per_part} or {b*self.num_query_groups_per_partition} or {b}"
+ )
+ else:
+ raise RuntimeError(f"DotProductAttention: Unsupported key.dim()={key.dim()}")
# [b, np, sq, sk]
- output_size = (query.size(1), query.size(2), query.size(0), key.size(0))
-
- # [sq, b, np, hn] -> [sq, b * np, hn]
- # This will be a simple view when doing normal attention, but in group query attention
- # the key and value tensors are repeated to match the queries so you can't use
- # simple strides to extract the queries.
- query = query.reshape(output_size[2], output_size[0] * output_size[1], -1)
- # [sk, b, np, hn] -> [sk, b * np, hn]
- key = key.view(output_size[3], output_size[0] * output_size[1], -1)
+ output_size = (b, np_per_part, sq_len, sk_len)
+
+ # Ensure value matches the [sk, b*np, hn] layout prior to context matmul
+ if value.dim() == 4:
+ # [sk, b, np, hn] -> [sk, b*np, hn]
+ value = value.reshape(value.size(0), b * np_per_part, -1)
+ elif value.dim() == 3:
+ if value.size(1) != b * np_per_part:
+ g = self.num_query_groups_per_partition
+ if value.size(1) == b * g:
+ r = value.size(2) // hn
+ if (r * g) == np_per_part and (value.size(2) % hn) == 0:
+ # [sk, b, g, r, hv] -> [sk, b, np, hv] -> [sk, b*np, hv]
+ value = (
+ value.reshape(value.size(0), b, g, r, hn)
+ .permute(0, 1, 3, 2, 4)
+ .reshape(value.size(0), b * np_per_part, hn)
+ )
+ elif np_per_part % g == 0 and r == 1:
+ repeat_times = np_per_part // g
+ value = value.repeat_interleave(repeat_times, dim=1)
+ else:
+ raise RuntimeError(
+ f"DotProductAttention: Cannot map value of shape {tuple(value.size())} to [sk, b*np, hv]. "
+ f"b={b}, np={np_per_part}, g={g}, hv={hn}"
+ )
+ elif value.size(1) == b:
+ value = value.repeat_interleave(np_per_part, dim=1)
+ else:
+ raise RuntimeError(
+ f"DotProductAttention: Unexpected value.size(1)={value.size(1)}; "
+ f"expected {b*np_per_part} or {b*self.num_query_groups_per_partition} or {b}"
+ )
+ else:
+ raise RuntimeError(f"DotProductAttention: Unsupported value.dim()={value.dim()}")
# preallocting input tensor: [b * np, sq, sk]
matmul_input_buffer = parallel_state.get_global_memory_buffer().get_tensor(
@@ -182,25 +349,39 @@ def forward(
# [sk, b, np, hn] --> [b, np, sq, hn]
# context layer shape: [b, np, sq, hn]
- output_size = (value.size(1), value.size(2), query.size(0), value.size(3))
+ # output_size = (value.size(1), value.size(2), query.size(0), value.size(3))
+ # Derive batch and heads-per-partition from attention_probs so this code
+ # works whether `value` was provided as [sk, b, np, hn] (4D) or
+ # [sk, b*np, hn] (3D). Use the last dimension of `value` for head dim.
+ b, np_per_part, sq_len, _sk_len = attention_probs.size()
+ head_dim = value.size(-1)
+ output_size = (b, np_per_part, sq_len, head_dim)
# change view [sk, b * np, hn]
- value = value.view(value.size(0), output_size[0] * output_size[1], -1)
+ # value = value.view(value.size(0), output_size[0] * output_size[1], -1)
+ value = value.reshape(value.size(0), output_size[0] * output_size[1], -1)
# change view [b * np, sq, sk]
- attention_probs = attention_probs.view(output_size[0] * output_size[1], output_size[2], -1)
+ # attention_probs = attention_probs.view(output_size[0] * output_size[1], output_size[2], -1)
+ attention_probs = attention_probs.reshape(output_size[0] * output_size[1], output_size[2], -1)
# matmul: [b * np, sq, hn]
context = torch.bmm(attention_probs, value.transpose(0, 1))
# change view [b, np, sq, hn]
- context = context.view(*output_size)
+ # context = context.view(*output_size)
+ # Use inferred head dimension to handle cases where `value` was packed differently
+ # and the product of dims differs from the expected head_dim.
+ context = context.reshape(output_size[0], output_size[1], output_size[2], -1)
# [b, np, sq, hn] --> [sq, b, np, hn]
context = context.permute(2, 0, 1, 3).contiguous()
# [sq, b, np, hn] --> [sq, b, hp]
new_context_shape = context.size()[:-2] + (self.hidden_size_per_partition,)
- context = context.view(*new_context_shape)
+ # context = context.view(*new_context_shape)
+ # Flatten heads using the actual last dimension to avoid shape mismatches when
+ # self.hidden_size_per_partition does not equal np_per_part * head_dim.
+ context = context.reshape(context.size(0), context.size(1), -1)
return context
diff --git a/aiak_megatron/megatron/core/transformer/dot_product_attention_no_te.py b/aiak_megatron/megatron/core/transformer/dot_product_attention_no_te.py
new file mode 100644
index 00000000..9374eee8
--- /dev/null
+++ b/aiak_megatron/megatron/core/transformer/dot_product_attention_no_te.py
@@ -0,0 +1,169 @@
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+
+import math
+from typing import Optional
+
+import torch
+from torch import Tensor
+
+from megatron.core import parallel_state, tensor_parallel
+from megatron.core.packed_seq_params import PackedSeqParams
+from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.utils import divide
+
+
+class ScaleMaskSoftmaxFallback:
+ """
+ Pure-PyTorch scale + mask + softmax used when fused kernels are unavailable.
+ """
+
+ def __init__(self, softmax_in_fp32: bool = False, attn_mask_type=None, scale: Optional[float] = None):
+ self.softmax_in_fp32 = softmax_in_fp32
+ self.attn_mask_type = attn_mask_type
+ self.scale = scale
+
+ def __call__(self, attention_scores: torch.Tensor, attention_mask: Optional[torch.Tensor], attn_mask_type: Optional[object] = None):
+ if self.scale is not None:
+ attention_scores = attention_scores * float(self.scale)
+
+ if attention_mask is not None:
+ if attention_mask.dtype == torch.bool:
+ add_mask = (~attention_mask).to(attention_scores.dtype) * -1e9
+ attention_scores = attention_scores + add_mask
+ else:
+ attention_scores = attention_scores + attention_mask
+
+ if self.softmax_in_fp32:
+ orig_dtype = attention_scores.dtype
+ probs = torch.softmax(attention_scores.float(), dim=-1).to(orig_dtype)
+ else:
+ probs = torch.softmax(attention_scores, dim=-1)
+
+ return probs
+
+
+class DotProductAttentionNoTE(MegatronModule):
+ """
+ Pure-PyTorch dot-product attention that avoids transformer-engine / fused softmax.
+ Shapes follow Megatron's convention: query/key/value are expected as [sq/sk, b, ng/np, hn].
+ """
+
+ def __init__(
+ self,
+ config: TransformerConfig,
+ layer_number: int,
+ attn_mask_type: AttnMaskType,
+ attention_type: str,
+ attention_dropout: float = None,
+ softmax_scale: float = None,
+ cp_comm_type: str = None,
+ ):
+ super().__init__(config=config)
+
+ self.config: TransformerConfig = config
+
+ assert self.config.context_parallel_size == 1, "Context parallelism is only supported by TEDotProductAttention!"
+ assert self.config.window_size is None, "Sliding Window Attention is only supported by TEDotProductAttention!"
+
+ self.layer_number = max(1, layer_number)
+ self.attn_mask_type = attn_mask_type
+ self.attention_type = attention_type
+
+ projection_size = self.config.kv_channels * self.config.num_attention_heads
+
+ world_size = parallel_state.get_tensor_model_parallel_world_size()
+ self.hidden_size_per_partition = divide(projection_size, world_size)
+ self.hidden_size_per_attention_head = divide(projection_size, config.num_attention_heads)
+ self.num_attention_heads_per_partition = divide(self.config.num_attention_heads, world_size)
+ self.num_query_groups_per_partition = divide(self.config.num_query_groups, world_size)
+
+ coeff = None
+ if softmax_scale is None:
+ self.softmax_scale = 1.0 / math.sqrt(self.hidden_size_per_attention_head)
+ else:
+ self.softmax_scale = softmax_scale
+
+ if self.config.apply_query_key_layer_scaling:
+ coeff = self.layer_number
+ self.softmax_scale /= coeff
+
+ self.scale_mask_softmax = ScaleMaskSoftmaxFallback(
+ softmax_in_fp32=self.config.attention_softmax_in_fp32,
+ attn_mask_type=self.attn_mask_type,
+ scale=coeff,
+ )
+
+ self.attention_dropout = torch.nn.Dropout(
+ self.config.attention_dropout if attention_dropout is None else attention_dropout
+ )
+
+ def forward(
+ self,
+ query: Tensor,
+ key: Tensor,
+ value: Tensor,
+ attention_mask: Tensor,
+ attn_mask_type: AttnMaskType = None,
+ attention_bias: Tensor = None,
+ packed_seq_params: Optional[PackedSeqParams] = None,
+ ):
+ packed_seq_params = None
+ assert packed_seq_params is None, (
+ "Packed sequence is not supported by DotProductAttentionNoTE."
+ "Please use TEDotProductAttention instead."
+ )
+ assert attention_bias is None, "Attention bias is not supported for DotProductAttentionNoTE."
+
+ if self.num_attention_heads_per_partition // self.num_query_groups_per_partition > 1:
+ key = key.repeat_interleave(
+ self.num_attention_heads_per_partition // self.num_query_groups_per_partition, dim=2
+ )
+ value = value.repeat_interleave(
+ self.num_attention_heads_per_partition // self.num_query_groups_per_partition, dim=2
+ )
+
+ output_size = (query.size(1), query.size(2), query.size(0), key.size(0))
+
+ query = query.reshape(output_size[2], output_size[0] * output_size[1], -1)
+ key = key.reshape(output_size[3], output_size[0] * output_size[1], -1)
+
+ matmul_input_buffer = parallel_state.get_global_memory_buffer().get_tensor(
+ (output_size[0] * output_size[1], output_size[2], output_size[3]), query.dtype, "mpu"
+ )
+
+ matmul_result = torch.baddbmm(
+ matmul_input_buffer,
+ query.transpose(0, 1),
+ key.transpose(0, 1).transpose(1, 2),
+ beta=0.0,
+ alpha=self.softmax_scale,
+ )
+
+ attention_scores = matmul_result.view(*output_size)
+
+ attention_probs: Tensor = self.scale_mask_softmax(attention_scores, attention_mask, attn_mask_type)
+
+ if not self.config.sequence_parallel:
+ with tensor_parallel.get_cuda_rng_tracker().fork():
+ attention_probs = self.attention_dropout(attention_probs)
+ else:
+ attention_probs = self.attention_dropout(attention_probs)
+
+ b, np_per_part, sq_len, _sk_len = attention_probs.size()
+ head_dim = value.size(-1)
+ output_size = (b, np_per_part, sq_len, head_dim)
+
+ value = value.reshape(value.size(0), output_size[0] * output_size[1], -1)
+ attention_probs = attention_probs.reshape(output_size[0] * output_size[1], output_size[2], -1)
+
+ context = torch.bmm(attention_probs, value.transpose(0, 1))
+
+ context = context.reshape(*output_size)
+ context = context.permute(2, 0, 1, 3).contiguous()
+
+ new_context_shape = context.size()[:-2] + (self.hidden_size_per_partition,)
+ context = context.reshape(*new_context_shape)
+
+ return context
diff --git a/aiak_megatron/megatron/core/transformer/moe/fused_a2a.py b/aiak_megatron/megatron/core/transformer/moe/fused_a2a.py
index 9bcbe8ff..64bd0724 100644
--- a/aiak_megatron/megatron/core/transformer/moe/fused_a2a.py
+++ b/aiak_megatron/megatron/core/transformer/moe/fused_a2a.py
@@ -17,9 +17,17 @@
HAVE_DEEP_EP = False
import torch
-from transformer_engine.pytorch.constants import TE_DType
-
+# from transformer_engine.pytorch.constants import TE_DType
+# Try to import TE_DType from Transformer Engine, but fall back if unavailable
logger = logging.getLogger(__name__)
+try:
+ from transformer_engine.pytorch.constants import TE_DType
+except Exception as e:
+ logger.warning(
+ f"TE_DType import failed, FP8 communication will be disabled: {str(e)}"
+ )
+ TE_DType = {}
+
try:
from transformer_engine.pytorch.tensor.float8_blockwise_tensor import (
Float8BlockQuantizer,
diff --git a/aiak_megatron/megatron/legacy/fused_kernels/__init__.py b/aiak_megatron/megatron/legacy/fused_kernels/__init__.py
index 74592952..a9f77ba9 100644
--- a/aiak_megatron/megatron/legacy/fused_kernels/__init__.py
+++ b/aiak_megatron/megatron/legacy/fused_kernels/__init__.py
@@ -15,6 +15,12 @@
def load(args):
+
+ # Skip fused kernels on ROCm (AMD GPUs)
+ import torch
+ if torch.version.hip is not None:
+ print("Skipping fused kernels compilation on ROCm/AMD platform")
+ return
# Check if cuda 11 is installed for compute capability 8.0
cc_flag = []
@@ -56,8 +62,20 @@ def _cpp_extention_load_helper(name, sources, extra_cuda_flags):
def _get_cuda_bare_metal_version(cuda_dir):
+ # Handle case where cuda_dir might be None or nvcc not in expected location
+ if cuda_dir is None:
+ cuda_dir = "/usr/local/cuda"
+
+ nvcc_path = cuda_dir + "/bin/nvcc"
+ if not os.path.exists(nvcc_path):
+ # Try alternate locations
+ for alt_path in ["/usr/local/cuda/bin/nvcc", "/usr/local/cuda-12.8/bin/nvcc", "/usr/local/cuda-12.1/bin/nvcc", "/usr/bin/nvcc"]:
+ if os.path.exists(alt_path):
+ nvcc_path = alt_path
+ break
+
raw_output = subprocess.check_output(
- [cuda_dir + "/bin/nvcc", "-V"], universal_newlines=True
+ [nvcc_path, "-V"], universal_newlines=True
)
output = raw_output.split()
release_idx = output.index("release") + 1
diff --git a/aiak_megatron/megatron/training/arguments.py b/aiak_megatron/megatron/training/arguments.py
index ac714238..0a4296fc 100644
--- a/aiak_megatron/megatron/training/arguments.py
+++ b/aiak_megatron/megatron/training/arguments.py
@@ -254,7 +254,9 @@ def validate_args(args, defaults={}):
)
if args.attention_backend == AttnBackend.local:
- assert args.spec[0] == 'local' , '--attention-backend local is only supported with --spec local'
+ # assert args.spec[0] == 'local' , '--attention-backend local is only supported with --spec local'
+ if args.spec is not None and len(args.spec) > 0:
+ assert args.spec[0] == 'local' , '--attention-backend local is only supported with --spec local'
# Pipeline model parallel size.
args.transformer_pipeline_model_parallel_size = args.pipeline_model_parallel_size
diff --git a/aiak_megatron/megatron/training/checkpointing.py b/aiak_megatron/megatron/training/checkpointing.py
index 4d89f0de..f78c8e6c 100644
--- a/aiak_megatron/megatron/training/checkpointing.py
+++ b/aiak_megatron/megatron/training/checkpointing.py
@@ -1283,14 +1283,214 @@ def load_checkpoint(model, optimizer, opt_param_scheduler, load_arg='load', stri
print_rank_0('could not find arguments in the checkpoint ...')
# Model.
- strict = False if args.retro_add_retriever else strict
+ # For fine-tuning/stage1 alignment, use non-strict loading to allow architecture mismatches.
+ # This is especially important when loading HF checkpoints into Megatron models.
+ if args.finetune or hasattr(args, 'trainable_modules'):
+ strict = False
+ else:
+ strict = False if args.retro_add_retriever else strict
if not skip_load_to_model_and_opt:
+ # Get device from model parameters
+ device = next(ddp_model[0].parameters()).device
+ print_rank_0(f"Loading checkpoint with device: {device}, strict={strict}")
+
+ # Move all tensors in state_dict to device before loading
+ model_state = {}
+ for k, v in state_dict['model'].items():
+ if isinstance(v, torch.Tensor):
+ model_state[k] = v.to(device)
+ else:
+ model_state[k] = v
+
if len(ddp_model) == 1:
- ddp_model[0].load_state_dict(state_dict['model'], strict=strict)
+ ddp_model[0].load_state_dict(model_state, strict=strict)
+ # Ensure entire model is on device (handles newly created layers from strict=False)
+ ddp_model[0].to(device)
+ # After checkpoint load, initialize any lazy TE linear modules with their loaded weights
+ try:
+ from megatron.core.extensions import transformer_engine as _te
+
+ for name, module in ddp_model[0].named_modules():
+ if isinstance(module, getattr(_te, '_FallbackLinear', tuple())):
+ # If this is a lazy linear (self.linear is None), try to initialize it from loaded weights
+ if hasattr(module, 'linear') and module.linear is None:
+ # Look for corresponding weights in the loaded state_dict
+ weight_key = None
+ bias_key = None
+ for state_key in model_state.keys():
+ if name in state_key and 'weight' in state_key:
+ weight_key = state_key
+ if name in state_key and 'bias' in state_key:
+ bias_key = state_key
+
+ if weight_key and weight_key in model_state:
+ weight = model_state[weight_key]
+ bias_tensor = model_state.get(bias_key) if bias_key else None
+ # Create linear with loaded shape
+ out_features, in_features = weight.shape
+ module.linear = torch.nn.Linear(in_features, out_features, bias=(bias_tensor is not None))
+ module.linear.weight.data.copy_(weight)
+ if bias_tensor is not None:
+ module.linear.bias.data.copy_(bias_tensor)
+ module.linear = module.linear.to(device=device, dtype=weight.dtype)
+ except Exception:
+ pass
+ # Ensure model params/buffers have correct dtype for mixed precision
+ try:
+ if getattr(args, 'bf16', False):
+ target_dtype = torch.bfloat16
+ elif getattr(args, 'fp16', False):
+ target_dtype = torch.float16
+ else:
+ target_dtype = None
+
+ if target_dtype is not None:
+ for p in ddp_model[0].parameters():
+ p.data = p.data.to(dtype=target_dtype, device=p.device)
+ for name, buf in ddp_model[0].named_buffers(recurse=True):
+ if buf is not None:
+ try:
+ buf.data = buf.data.to(dtype=target_dtype, device=buf.device)
+ except Exception:
+ # some buffers may be non-tensor or empty
+ pass
+ except Exception:
+ pass
+ # Also ensure any fallback TE linear modules are moved to device
+ try:
+ from megatron.core.extensions import transformer_engine as _te
+
+ for m in ddp_model[0].modules():
+ if isinstance(m, getattr(_te, '_FallbackLinear', tuple())):
+ m.to(device)
+ except Exception:
+ pass
+ # Final dtype enforcement after moving any fallback modules
+ try:
+ if getattr(args, 'bf16', False):
+ target_dtype = torch.bfloat16
+ elif getattr(args, 'fp16', False):
+ target_dtype = torch.float16
+ else:
+ target_dtype = None
+
+ if target_dtype is not None:
+ mismatches = []
+ total = 0
+ for name, p in ddp_model[0].named_parameters():
+ total += 1
+ if p.dtype != target_dtype:
+ mismatches.append((name, str(p.device), str(p.dtype)))
+ p.data = p.data.to(device=p.device, dtype=target_dtype)
+ for name, b in ddp_model[0].named_buffers(recurse=True):
+ if isinstance(b, torch.Tensor) and b.dtype != target_dtype:
+ mismatches.append((f'buffer:{name}', str(b.device), str(b.dtype)))
+ try:
+ b.data = b.data.to(device=b.device, dtype=target_dtype)
+ except Exception:
+ pass
+ if mismatches:
+ print_rank_0(f'Post-load converted {len(mismatches)}/{total} params/buffers to {target_dtype}. Examples: {mismatches[:10]}')
+ except Exception:
+ pass
else:
for i in range(len(ddp_model)):
mpu.set_virtual_pipeline_model_parallel_rank(i)
- ddp_model[i].load_state_dict(state_dict['model%d' % i], strict=strict)
+ model_state_i = {}
+ for k, v in state_dict['model%d' % i].items():
+ if isinstance(v, torch.Tensor):
+ model_state_i[k] = v.to(device)
+ else:
+ model_state_i[k] = v
+ ddp_model[i].load_state_dict(model_state_i, strict=strict)
+ # Ensure entire model is on device
+ ddp_model[i].to(device)
+ # After checkpoint load, initialize any lazy TE linear modules with their loaded weights
+ try:
+ from megatron.core.extensions import transformer_engine as _te
+
+ for name, module in ddp_model[i].named_modules():
+ if isinstance(module, getattr(_te, '_FallbackLinear', tuple())):
+ # If this is a lazy linear (self.linear is None), try to initialize it from loaded weights
+ if hasattr(module, 'linear') and module.linear is None:
+ # Look for corresponding weights in the loaded state_dict
+ weight_key = None
+ bias_key = None
+ for state_key in model_state_i.keys():
+ if name in state_key and 'weight' in state_key:
+ weight_key = state_key
+ if name in state_key and 'bias' in state_key:
+ bias_key = state_key
+
+ if weight_key and weight_key in model_state_i:
+ weight = model_state_i[weight_key]
+ bias_tensor = model_state_i.get(bias_key) if bias_key else None
+ # Create linear with loaded shape
+ out_features, in_features = weight.shape
+ module.linear = torch.nn.Linear(in_features, out_features, bias=(bias_tensor is not None))
+ module.linear.weight.data.copy_(weight)
+ if bias_tensor is not None:
+ module.linear.bias.data.copy_(bias_tensor)
+ module.linear = module.linear.to(device=device, dtype=weight.dtype)
+ except Exception:
+ pass
+ # Ensure any fallback TE linears are on the same device
+ try:
+ from megatron.core.extensions import transformer_engine as _te
+
+ for m in ddp_model[i].modules():
+ if isinstance(m, getattr(_te, '_FallbackLinear', tuple())):
+ m.to(device)
+ except Exception:
+ pass
+ # Ensure model params/buffers have correct dtype for mixed precision
+ try:
+ if getattr(args, 'bf16', False):
+ target_dtype = torch.bfloat16
+ elif getattr(args, 'fp16', False):
+ target_dtype = torch.float16
+ else:
+ target_dtype = None
+
+ if target_dtype is not None:
+ for p in ddp_model[i].parameters():
+ p.data = p.data.to(dtype=target_dtype, device=p.device)
+ for name, buf in ddp_model[i].named_buffers(recurse=True):
+ if buf is not None:
+ try:
+ buf.data = buf.data.to(dtype=target_dtype, device=buf.device)
+ except Exception:
+ pass
+ except Exception:
+ pass
+ # Final dtype enforcement after moving any fallback modules
+ try:
+ if getattr(args, 'bf16', False):
+ target_dtype = torch.bfloat16
+ elif getattr(args, 'fp16', False):
+ target_dtype = torch.float16
+ else:
+ target_dtype = None
+
+ if target_dtype is not None:
+ mismatches = []
+ total = 0
+ for name, p in ddp_model[i].named_parameters():
+ total += 1
+ if p.dtype != target_dtype:
+ mismatches.append((name, str(p.device), str(p.dtype)))
+ p.data = p.data.to(device=p.device, dtype=target_dtype)
+ for name, b in ddp_model[i].named_buffers(recurse=True):
+ if isinstance(b, torch.Tensor) and b.dtype != target_dtype:
+ mismatches.append((f'buffer:{name}', str(b.device), str(b.dtype)))
+ try:
+ b.data = b.data.to(device=b.device, dtype=target_dtype)
+ except Exception:
+ pass
+ if mismatches:
+ print_rank_0(f'Post-load converted {len(mismatches)}/{total} params/buffers to {target_dtype}. Examples: {mismatches[:10]}')
+ except Exception:
+ pass
# Fix up query/key/value matrix ordering if needed.
checkpoint_version = get_checkpoint_version()
diff --git a/aiak_megatron/megatron/training/checkpointing_org.py b/aiak_megatron/megatron/training/checkpointing_org.py
new file mode 100644
index 00000000..033e2bb4
--- /dev/null
+++ b/aiak_megatron/megatron/training/checkpointing_org.py
@@ -0,0 +1,1444 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+"""Input/output checkpointing."""
+
+import contextlib
+import os
+import random
+import shutil
+import sys
+import threading
+from enum import Enum, auto
+from logging import getLogger
+from pathlib import Path
+
+import numpy as np
+from time import time
+
+import torch
+
+from megatron.core import mpu, tensor_parallel, dist_checkpointing
+from megatron.core.dist_checkpointing.mapping import ShardedObject
+from megatron.core.dist_checkpointing.serialization import get_default_load_sharded_strategy
+from megatron.core.dist_checkpointing.strategies.fully_parallel import \
+ FullyParallelSaveStrategyWrapper, FullyParallelLoadStrategyWrapper
+from megatron.core.num_microbatches_calculator import update_num_microbatches
+from megatron.core.fp8_utils import is_float8tensor, dequantize_fp8_tensor
+from megatron.core.rerun_state_machine import get_rerun_state_machine
+from .async_utils import schedule_async_save, is_empty_async_queue
+from .global_vars import get_args, get_one_logger
+from .utils import unwrap_model, print_rank_0, append_to_progress_log, is_last_rank
+from ..core.dist_checkpointing.serialization import \
+ get_default_save_sharded_strategy
+from .one_logger_utils import on_save_checkpoint_start, on_save_checkpoint_success
+from . import wandb_utils
+from . import ft_integration
+
+# [ModelOpt]: Import
+try:
+ from modelopt.torch.opt.plugins import (
+ save_modelopt_state,
+ save_sharded_modelopt_state,
+ restore_modelopt_state,
+ restore_sharded_modelopt_state,
+ )
+ has_nvidia_modelopt = True
+except Exception:
+ has_nvidia_modelopt = False
+
+_CHECKPOINT_VERSION = None
+
+logger = getLogger(__name__)
+_NON_PERSISTENT_CKPT_SUBDIR = 'non_persistent'
+
+
+def set_checkpoint_version(value):
+ """set the current checkpoint version"""
+ global _CHECKPOINT_VERSION
+ if _CHECKPOINT_VERSION is not None:
+ assert _CHECKPOINT_VERSION == value, \
+ "checkpoint versions do not match"
+ _CHECKPOINT_VERSION = value
+
+
+def get_checkpoint_version():
+ """get the current checkpoint version"""
+ global _CHECKPOINT_VERSION
+ return _CHECKPOINT_VERSION
+
+
+def check_checkpoint_args(checkpoint_args):
+ """Ensure fixed arguments for a model are the same for the input
+ arguments and the one retrieved from checkpoint."""
+ args = get_args()
+
+ def _compare(arg_name, old_arg_name=None, default=None):
+ if old_arg_name is not None:
+ ckpt_arg_name = old_arg_name
+ else:
+ ckpt_arg_name = arg_name
+ if default is not None:
+ checkpoint_value = getattr(checkpoint_args, ckpt_arg_name, default)
+ else:
+ checkpoint_value = getattr(checkpoint_args, ckpt_arg_name)
+ args_value = getattr(args, arg_name)
+ error_message = '{} value from checkpoint ({}) is not equal to the ' \
+ 'input argument value ({}).'.format(
+ arg_name, checkpoint_value, args_value)
+ assert checkpoint_value == args_value, error_message
+
+ _compare('num_layers')
+ _compare('hidden_size')
+ _compare('num_attention_heads')
+ _compare('add_position_embedding', default=False)
+ if args.vocab_file:
+ _compare('max_position_embeddings')
+ _compare('make_vocab_size_divisible_by')
+ if not args.use_dist_ckpt:
+ _compare('padded_vocab_size')
+ _compare('tokenizer_type')
+ if args.data_parallel_random_init:
+ _compare('data_parallel_random_init')
+ if get_checkpoint_version() < 3.0:
+ _compare('tensor_model_parallel_size',
+ old_arg_name='model_parallel_size')
+ if get_checkpoint_version() >= 3.0 and not args.use_dist_ckpt:
+ _compare('tensor_model_parallel_size')
+ _compare('pipeline_model_parallel_size')
+
+
+def ensure_directory_exists(filename, check_parent=True):
+ """Build filename's path if it does not already exists."""
+ dirname = os.path.dirname(filename) if check_parent else filename
+ os.makedirs(dirname, exist_ok=True)
+
+
+def get_checkpoint_name(checkpoints_path, iteration, release=False,
+ pipeline_parallel=None,
+ tensor_rank=None, pipeline_rank=None,
+ expert_parallel=None, expert_rank=None,
+ return_base_dir=False,
+ basename="model_optim_rng.pt"):
+ """Determine the directory name for this rank's checkpoint."""
+ if release:
+ directory = 'release'
+ else:
+ directory = 'iter_{:07d}'.format(iteration)
+ if return_base_dir:
+ common_path = os.path.join(checkpoints_path, directory)
+ return common_path
+
+ # Use both the tensor and pipeline MP rank.
+ if pipeline_parallel is None:
+ pipeline_parallel = (mpu.get_pipeline_model_parallel_world_size() > 1)
+ if tensor_rank is None:
+ tensor_rank = mpu.get_tensor_model_parallel_rank()
+ if pipeline_rank is None:
+ pipeline_rank = mpu.get_pipeline_model_parallel_rank()
+ if expert_parallel is None:
+ expert_parallel = (mpu.get_expert_model_parallel_world_size() > 1)
+ if expert_rank is None:
+ expert_rank = mpu.get_expert_model_parallel_rank()
+
+ # Use both the tensor and pipeline MP rank. If using the distributed
+ # optimizer, then the optimizer's path must additionally include the
+ # data parallel rank.
+ if not pipeline_parallel:
+ common_path = os.path.join(checkpoints_path, directory,
+ f'mp_rank_{tensor_rank:02d}')
+ else:
+ common_path = os.path.join(checkpoints_path, directory,
+ f'mp_rank_{tensor_rank:02d}_{pipeline_rank:03d}')
+
+ if expert_parallel:
+ common_path = common_path + f'_{expert_rank:03d}'
+
+ return os.path.join(common_path, basename)
+
+
+def get_distributed_optimizer_checkpoint_name(model_checkpoint_name):
+ return os.path.join(os.path.dirname(model_checkpoint_name),
+ "distrib_optim.pt")
+
+
+def find_checkpoint_rank_0(checkpoints_path, iteration, release=False):
+ """Finds the checkpoint for rank 0 without knowing if we are using
+ pipeline parallelism/expert parallelism or not.
+
+ Since the checkpoint naming scheme changes if pipeline or expert
+ parallelism is present, we need to look for both naming schemes if
+ we don't know if the checkpoint has pipeline or expert parallelism.
+ """
+
+ # Look for checkpoint with no pipelining and no expert parallelism
+ filename = get_checkpoint_name(checkpoints_path, iteration, release,
+ pipeline_parallel=False,
+ tensor_rank=0, pipeline_rank=0,
+ expert_parallel=False, expert_rank=0)
+ if os.path.isfile(filename):
+ return filename
+
+ # Look for checkpoint with no pipelining and expert parallelism
+ filename = get_checkpoint_name(checkpoints_path, iteration, release,
+ pipeline_parallel=False,
+ tensor_rank=0, pipeline_rank=0,
+ expert_parallel=True, expert_rank=0)
+ if os.path.isfile(filename):
+ return filename
+
+ # Look for checkpoint with pipelining and no expert parallelism
+ filename = get_checkpoint_name(checkpoints_path, iteration, release,
+ pipeline_parallel=True,
+ tensor_rank=0, pipeline_rank=0,
+ expert_parallel=False, expert_rank=0)
+ if os.path.isfile(filename):
+ return filename
+
+ # Look for checkpoint with pipelining and expert parallelism
+ filename = get_checkpoint_name(checkpoints_path, iteration, release,
+ pipeline_parallel=True,
+ tensor_rank=0, pipeline_rank=0,
+ expert_parallel=True, expert_rank=0)
+ if os.path.isfile(filename):
+ return filename
+
+ # Look for a distributed checkpoint
+ filename = get_checkpoint_name(checkpoints_path, iteration, release,
+ pipeline_parallel=True,
+ return_base_dir=True)
+ if dist_checkpointing.check_is_distributed_checkpoint(filename):
+ return filename
+
+ return None
+
+
+def get_checkpoint_tracker_filename(checkpoints_path):
+
+ """Tracker file rescords the latest chckpoint during
+ training to restart from."""
+ return os.path.join(checkpoints_path, 'latest_checkpointed_iteration.txt')
+
+
+def checkpoint_exists(checkpoints_path):
+ """Check if a checkpoint exists."""
+ if checkpoints_path is None:
+ return False
+ load_step = 'latest_checkpointed_iteration.txt'
+ return os.path.exists(os.path.join(checkpoints_path, load_step))
+
+
+def read_metadata(tracker_filename):
+ # Read the tracker file and either set the iteration or
+ # mark it as a release checkpoint.
+ iteration = 0
+ release = False
+ with open(tracker_filename, 'r') as f:
+ metastring = f.read().strip()
+ try:
+ iteration = int(metastring)
+ except ValueError:
+ release = metastring == 'release'
+ if not release:
+ print_rank_0('ERROR: Invalid metadata file {}. Exiting'.format(
+ tracker_filename))
+ sys.exit()
+ assert iteration > 0 or release, 'error parsing metadata file {}'.format(
+ tracker_filename)
+
+ # Get the max iteration retrieved across the ranks.
+ if torch.distributed.is_initialized():
+ iters_cuda = torch.tensor([iteration], dtype=torch.long, device='cuda')
+ torch.distributed.all_reduce(iters_cuda, op=torch.distributed.ReduceOp.MAX)
+ max_iter = iters_cuda[0].item()
+
+ # We should now have all the same iteration.
+ # If not, print a warning and chose the maximum
+ # iteration across all ranks.
+ if iteration != max_iter:
+ rank = torch.distributed.get_rank()
+ print('WARNING: on rank {} found iteration {} in the '
+ 'metadata while max iteration across the ranks '
+ 'is {}, replacing it with max iteration.'.format(
+ rank, iteration, max_iter), flush=True)
+ else:
+ # When loading a checkpoint outside of training (for example,
+ # when editing it), we might not have torch distributed
+ # initialized, in this case, just assume we have the latest
+ max_iter = iteration
+ return max_iter, release
+
+
+def get_rng_state(use_dist_ckpt: bool = False):
+ """ collect rng state across data parallel ranks """
+ args = get_args()
+ rng_state = {
+ 'random_rng_state': random.getstate(),
+ 'np_rng_state': np.random.get_state(),
+ 'torch_rng_state': torch.get_rng_state(),
+ 'cuda_rng_state': torch.cuda.get_rng_state(),
+ 'rng_tracker_states': tensor_parallel.get_cuda_rng_tracker().get_states()}
+
+ rng_state_list = None
+ if torch.distributed.is_initialized() and \
+ mpu.get_data_parallel_world_size() > 1 and \
+ args.data_parallel_random_init:
+ rng_state_list = \
+ [None for i in range(mpu.get_data_parallel_world_size())]
+ torch.distributed.all_gather_object(
+ rng_state_list,
+ rng_state,
+ group=mpu.get_data_parallel_group())
+ else:
+ rng_state_list = [rng_state]
+
+ if use_dist_ckpt:
+ pp_rank = mpu.get_pipeline_model_parallel_rank()
+ pp_size = mpu.get_pipeline_model_parallel_world_size()
+ tp_rank = mpu.get_tensor_model_parallel_rank()
+ tp_size = mpu.get_tensor_model_parallel_world_size()
+ rng_state_list = ShardedObject('rng_state', rng_state_list, (pp_size, tp_size), (pp_rank, tp_rank),
+ replica_id=mpu.get_data_parallel_rank(with_context_parallel=True))
+
+ return rng_state_list
+
+
+class CheckpointType(Enum):
+ """Checkpoint type"""
+ LEGACY = auto()
+ LOCAL = auto()
+ GLOBAL = auto()
+
+
+def save_checkpoint(iteration, model, optimizer, opt_param_scheduler, num_floating_point_operations_so_far,
+ checkpointing_context=None, save_arg='save', pipeline_rank=None, expert_rank=None,
+ tensor_rank=None, pipeline_parallel=None, expert_parallel=None, non_persistent_ckpt=False,
+ train_data_iterator=None, preprocess_common_state_dict_fn=None):
+ """Save a model, optimizer and optionally dataloader checkpoint.
+
+ Checkpointing context is used to persist some checkpointing state
+ throughout a single job. Must be initialized externally (not used if None).
+
+ If non_persistent_ckpt is True,
+ the checkpoint will be saved with special functionality for removing old checkpoints.
+ There are several types of non-persistent checkpoints:
+ "global" - Saved as a standard checkpoint (e.g., on Lustre) with old checkpoints being removed.
+ "local" - Each rank saves a portion of the checkpoint locally (e.g., on SSD/ramdisk).
+
+ Dataloader checkpoint is only saved if the dataloader supports it. Currently this applies only
+ to the Megatron Energon dataloader (multimodal) and not the built-in Megatron dataloader (text-only).
+ """
+ start_ckpt = time()
+ args = get_args()
+ save_dir = getattr(args, save_arg)
+
+ if args.async_save and not is_empty_async_queue():
+ print_rank_0('WARNING: Starting a checkpoint save before previous has finished. '
+ 'Consider increasing the checkpoint interval.')
+
+ # Prepare E2E metrics at start of save checkpoint
+ productive_metrics = on_save_checkpoint_start(args.async_save)
+
+ # Monitor for the checkpointing timeout (no-op if FT is not enabled)
+ ft_integration.on_checkpointing_start()
+
+ # Only rank zero of the data parallel writes to the disk.
+ model = unwrap_model(model)
+
+ # Handle non_persistent_ckpt flag. Besides overwriting `args.save` and
+ # `args.use_dist_ckpt`, non-persistent global ckpt requires no additional logic
+ ckpt_type = CheckpointType.GLOBAL if args.use_dist_ckpt else CheckpointType.LEGACY
+ if non_persistent_ckpt:
+ if args.non_persistent_ckpt_type == 'global':
+ ckpt_type = CheckpointType.GLOBAL
+ save_dir = (
+ args.non_persistent_global_ckpt_dir
+ if args.non_persistent_global_ckpt_dir
+ else os.path.join(save_dir, _NON_PERSISTENT_CKPT_SUBDIR)
+ )
+ # TODO Can we ensure the previous checkpoint is saved? We don't want to allow two saves in parallel.
+ cleanup_old_non_persistent_checkpoint(save_dir, leave_ckpt_num=1, do_async=args.async_save)
+ elif args.non_persistent_ckpt_type == 'local':
+ ckpt_type = CheckpointType.LOCAL
+ save_dir = checkpointing_context['local_checkpoint_manager'].local_ckpt_dir
+ else:
+ assert False, \
+ f'Please use local or global non-persistent checkpoints (got: {args.non_persistent_ckpt_type})'
+
+ ckpt_format = args.ckpt_format if ckpt_type == CheckpointType.GLOBAL else 'torch'
+ print_rank_0('saving checkpoint at iteration {:7d} to {} in {} format'.format(iteration, save_dir, ckpt_format))
+
+ # Collect rng state across data parallel ranks.
+ rng_state = get_rng_state(ckpt_type != CheckpointType.LEGACY)
+
+ # Collect rerun state across all ranks
+ rerun_state_machine = get_rerun_state_machine()
+ rerun_state = rerun_state_machine.state_dict(
+ data_iterator=train_data_iterator, use_dist_ckpt=ckpt_type != CheckpointType.LEGACY
+ )
+
+ # Checkpoint name.
+ return_base_dir = (ckpt_type != CheckpointType.LEGACY)
+ checkpoint_name = get_checkpoint_name(save_dir, iteration, release=False,
+ pipeline_parallel=pipeline_parallel, tensor_rank=tensor_rank,
+ pipeline_rank=pipeline_rank, expert_parallel=expert_parallel,
+ expert_rank=expert_rank, return_base_dir=return_base_dir)
+
+ # Save dataloader state if the dataloader supports it (currently only Megatron Energon).
+ maybe_save_dataloader_state(train_data_iterator, iteration, getattr(args, "dataloader_save", None))
+
+ # Save distributed optimizer's custom parameter state.
+ if (
+ args.use_distributed_optimizer
+ and not args.no_save_optim
+ and optimizer is not None
+ and ckpt_type == CheckpointType.LEGACY
+ ):
+ optim_checkpoint_name = \
+ get_distributed_optimizer_checkpoint_name(checkpoint_name)
+ ensure_directory_exists(optim_checkpoint_name)
+ if not optimizer.is_stub_optimizer:
+ optimizer.save_parameter_state(optim_checkpoint_name)
+
+ async_save_request = None
+ if args.async_save:
+ if ckpt_type == CheckpointType.LEGACY:
+ raise NotImplementedError('Async checkpoint save not implemented for legacy checkpoints')
+ elif ckpt_type == CheckpointType.GLOBAL and args.ckpt_format != 'torch_dist':
+ raise NotImplementedError(f'Async checkpoint save not implemented for {args.ckpt_format} '
+ f'distributed checkpoint format')
+
+ rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0
+
+ # Collect args, model, RNG.
+ if not torch.distributed.is_initialized() \
+ or mpu.get_expert_data_parallel_rank() == 0 \
+ or ckpt_type != CheckpointType.LEGACY:
+ optim_sd_kwargs = {}
+ if ckpt_type != CheckpointType.LEGACY and args.use_distributed_optimizer:
+ optim_sd_kwargs['sharding_type'] = ('fully_sharded_model_space'
+ if args.ckpt_fully_parallel_save
+ else 'dp_zero_gather_scatter')
+ print_rank_0(f'Storing distributed optimizer sharded state of type {optim_sd_kwargs["sharding_type"]}')
+
+ state_dict = generate_state_dict(
+ args,
+ model,
+ optimizer,
+ opt_param_scheduler,
+ rng_state,
+ use_dist_ckpt=ckpt_type != CheckpointType.LEGACY,
+ iteration=iteration,
+ optim_sd_kwargs=optim_sd_kwargs,
+ rerun_state=rerun_state,
+ )
+
+ state_dict['num_floating_point_operations_so_far'] = num_floating_point_operations_so_far
+ if ckpt_type == CheckpointType.GLOBAL:
+ if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0:
+ # TODO Handle non-empty directories (e.g., after a crash during saving).
+ ensure_directory_exists(checkpoint_name, check_parent=False)
+ if checkpointing_context is not None and 'save_strategy' in checkpointing_context:
+ save_strategy = checkpointing_context['save_strategy']
+ # Already saved once before - don't need to rerun sharding validation
+ validate_sharding_integrity = not args.ckpt_assume_constant_structure
+ else:
+ validate_sharding_integrity = True
+ save_strategy = get_default_save_sharded_strategy(args.ckpt_format)
+ if args.ckpt_assume_constant_structure and args.ckpt_format == 'torch_dist':
+ save_strategy.use_cached_ckpt_structure = args.ckpt_assume_constant_structure
+ if checkpointing_context is not None and 'load_strategy' in checkpointing_context:
+ cached_global_metadata = getattr(checkpointing_context['load_strategy'],
+ 'cached_global_metadata', None)
+ if cached_global_metadata is not None:
+ logger.debug("Plugging in the read metadata from the load strategy...")
+ save_strategy.cached_global_metadata = cached_global_metadata
+ else:
+ logger.debug("Failed to plug in the read metadata from the load strategy...")
+
+ if args.ckpt_fully_parallel_save:
+ save_strategy = FullyParallelSaveStrategyWrapper(
+ save_strategy,
+ mpu.get_data_parallel_group(with_context_parallel=True),
+ args.ckpt_assume_constant_structure
+ )
+ # Store save strategy for future checkpoint saves
+ if checkpointing_context is not None:
+ checkpointing_context['save_strategy'] = save_strategy
+ end_ckpt = time()
+ logger.debug(f"rank: {rank}, takes {end_ckpt - start_ckpt} to prepare state dict for ckpt ")
+ async_save_request = dist_checkpointing.save(
+ state_dict, checkpoint_name, save_strategy,
+ async_sharded_save=args.async_save,
+ validate_access_integrity=validate_sharding_integrity,
+ preprocess_common_before_consistancy_check=preprocess_common_state_dict_fn)
+ # [ModelOpt]: save sharded modelopt_state
+ if has_nvidia_modelopt:
+ save_sharded_modelopt_state(model, checkpoint_name, (args.ckpt_format, 1))
+ else:
+ # [ModelOpt]: Inject modelopt_state into state_dict
+ if has_nvidia_modelopt:
+ if ckpt_type == CheckpointType.LOCAL:
+ print_rank_0('WARNING: Local checkpointing does not support nvidia_modelopt.')
+ else:
+ save_modelopt_state(model, state_dict)
+
+ end_ckpt = time()
+ logger.debug(f"rank: {rank}, takes {end_ckpt - start_ckpt} to prepare state dict for ckpt ")
+ if ckpt_type == CheckpointType.LOCAL:
+ try:
+ from megatron.core.dist_checkpointing.tensor_aware_state_dict import MCoreTensorAwareStateDict
+ except ModuleNotFoundError:
+ raise RuntimeError("The 'nvidia_resiliency_ext' module is required for local "
+ "checkpointing but was not found. Please ensure it is installed.")
+
+ algo = args.non_persistent_local_ckpt_algo
+ cached_metadata = None
+ if args.ckpt_assume_constant_structure and 'local_checkpoint_cache' in checkpointing_context:
+ cached_metadata = checkpointing_context['local_checkpoint_cache']
+ state_dict_for_save, cacheable_metadata = MCoreTensorAwareStateDict.from_state_dict(
+ state_dict, algo=algo, cached_metadata=cached_metadata,
+ parallelization_group=mpu.get_data_parallel_group(with_context_parallel=True)
+ )
+ async_save_request = checkpointing_context['local_checkpoint_manager'].save(
+ state_dict_for_save, iteration, is_async=bool(args.async_save)
+ )
+ checkpointing_context['local_checkpoint_cache'] = cacheable_metadata
+ else:
+ assert ckpt_type == CheckpointType.LEGACY
+ # Save.
+ ensure_directory_exists(checkpoint_name)
+ torch.save(state_dict, checkpoint_name)
+ start_misc = time()
+ if ckpt_type != CheckpointType.LOCAL:
+ if not args.async_save:
+ assert async_save_request is None
+ # Wait so everyone is done (necessary)
+ if torch.distributed.is_initialized():
+ torch.distributed.barrier()
+
+ # And update the latest iteration
+ if not torch.distributed.is_initialized() \
+ or torch.distributed.get_rank() == 0:
+ tracker_filename = get_checkpoint_tracker_filename(save_dir)
+
+ if ckpt_type == CheckpointType.LOCAL:
+ def iter_finalize_fn():
+ print_rank_0(' successfully saved local checkpoint from iteration {:7d}'.format(iteration))
+ if args.log_progress and args.async_save:
+ append_to_progress_log(f'Saved async local checkpoint\tIteration: {iteration}', barrier=False)
+ else:
+ def iter_finalize_fn():
+ with open(tracker_filename, 'w') as f:
+ f.write(str(iteration))
+
+ _tensor_rank = tensor_rank if tensor_rank is not None else mpu.get_tensor_model_parallel_rank()
+ _pipeline_rank = pipeline_rank if pipeline_rank is not None else mpu.get_pipeline_model_parallel_rank()
+
+ print_rank_0(f' successfully saved checkpoint from iteration {int(iteration):7d} to {args.save} '
+ f'[ t {_tensor_rank + 1}/{mpu.get_tensor_model_parallel_world_size()}, '
+ f'p {_pipeline_rank + 1}/{mpu.get_pipeline_model_parallel_world_size()} ]')
+ if args.log_progress and args.async_save:
+ append_to_progress_log(f'Saved async checkpoint\tIteration: {iteration}', barrier=False)
+
+ if args.async_save:
+ assert async_save_request is not None
+ async_save_request.add_finalize_fn(iter_finalize_fn)
+ else:
+ iter_finalize_fn()
+
+ # Additional callback for one_logger (last rank)
+ if not torch.distributed.is_initialized() \
+ or is_last_rank():
+ def onelogger_finalize_fn():
+ on_save_checkpoint_success(productive_metrics, args.async_save)
+ if args.async_save:
+ assert async_save_request is not None
+ async_save_request.add_finalize_fn(onelogger_finalize_fn)
+ else:
+ onelogger_finalize_fn()
+
+ # Additional callback for wandb (last rank)
+ if not torch.distributed.is_initialized() \
+ or is_last_rank():
+ def wandb_finalize_fn():
+ wandb_utils.on_save_checkpoint_success(checkpoint_name,
+ get_checkpoint_tracker_filename(save_dir), save_dir, iteration)
+ if args.async_save:
+ assert async_save_request is not None
+ async_save_request.add_finalize_fn(wandb_finalize_fn)
+ else:
+ wandb_finalize_fn()
+
+ if args.async_save:
+ schedule_async_save(async_save_request)
+ print_rank_0(' scheduled an async checkpoint save at iteration {:7d} to {}' \
+ .format(iteration, save_dir))
+
+ # Wait so everyone is done (not necessary)
+ if torch.distributed.is_initialized():
+ torch.distributed.barrier()
+
+ end_misc = time()
+ logger.debug(f"rank: {rank}, takes {end_misc - start_misc} to finalize ckpt save ")
+
+ ft_integration.on_checkpointing_end(is_async_finalization=False)
+
+
+def cleanup_old_non_persistent_checkpoint(save_dir, leave_ckpt_num=1, do_async=False):
+ """Clean up old non-persistent checkpoints"""
+ if torch.distributed.is_initialized() and torch.distributed.get_rank() != 0:
+ return
+ save_dir = Path(save_dir)
+
+ iter_prefix = "iter_"
+ iter_ckpts = save_dir.rglob(f'{iter_prefix}*')
+ sorted_iter_ckpts = sorted(iter_ckpts, key=lambda ckpt_name: int(ckpt_name.name[len(iter_prefix):]))
+ if not sorted_iter_ckpts:
+ return
+ rm_iter_ckpts = sorted_iter_ckpts[:-leave_ckpt_num]
+ print_rank_0(f'Non-persistent checkpoints scheduled for removal: {rm_iter_ckpts}')
+ print_rank_0(f'Non-persistent checkpoints to be kept: {sorted_iter_ckpts[-leave_ckpt_num:]}')
+
+ def remove_iter_ckpts(_iter_ckpts):
+ for ckpt in _iter_ckpts:
+ shutil.rmtree(ckpt)
+ if do_async:
+ threading.Thread(target=remove_iter_ckpts, args=(rm_iter_ckpts,)).start()
+ else:
+ remove_iter_ckpts(rm_iter_ckpts)
+
+
+def maybe_save_dataloader_state(train_iterator, iteration, dataloader_save_path):
+ """Saves dataloader state if the dataloader supports it.
+
+ Currently, this is only used by Megatron Energon dataloader (multimodal) to store its state at a
+ specific iteration. The Megatron built-in dataloader (text-only) creates index files upfront
+ to track its state.
+
+ If the provided dataloader has `save_state` method, then it is called to save the state.
+ Otherwise, no state is saved.
+
+ Args:
+ train_iterator (iterable): Train dataloader.
+ iteration (int): Current iteration.
+ dataloader_save_path (str): Path where the dataloader state is saved.
+ """
+ # If no dataloader or saving path is provided, exit early, otherwise, raise an error.
+ if train_iterator is None or dataloader_save_path is None or dataloader_save_path == "":
+ return
+
+ # If dataloader doesn't support saving state, raise an error.
+ if not hasattr(train_iterator.iterable, "save_state"):
+ raise RuntimeError(f"Could not find a save_state for the train_iterator of type {type(train_iterator)}")
+
+ # Save dataloader state for each data parallel rank only once.
+ first_rank = mpu.is_pipeline_first_stage(ignore_virtual=True) and mpu.get_tensor_model_parallel_rank() == 0
+ if not first_rank:
+ return
+
+ dp_rank = mpu.get_data_parallel_rank()
+ print(f"saving dataloader checkpoint at iteration {iteration} to {dataloader_save_path}")
+ train_dataloader_state_dict = train_iterator.iterable.save_state()
+ data_state_save_path = get_checkpoint_name(
+ dataloader_save_path, iteration,
+ basename=f'train_dataloader_dprank{dp_rank:03d}.pt'
+ )
+
+ torch.distributed.barrier(group=mpu.get_data_parallel_group())
+
+ if mpu.get_data_parallel_rank() == 0:
+ ensure_directory_exists(data_state_save_path)
+
+ torch.distributed.barrier(group=mpu.get_data_parallel_group())
+
+ dataloader_save_dict = {}
+ dataloader_save_dict['dataloader_state_dict'] = train_dataloader_state_dict
+ torch.save(dataloader_save_dict, data_state_save_path)
+
+
+def generate_state_dict(args, model, optimizer, opt_param_scheduler,
+ rng_state, use_dist_ckpt=False, iteration=None,
+ optim_sd_kwargs=None, rerun_state=None):
+ # Arguments, iteration, and model.
+ state_dict = {}
+ state_dict['args'] = args
+ state_dict['checkpoint_version'] = 3.0
+ if iteration is not None:
+ state_dict['iteration'] = iteration
+
+ if len(model) == 1:
+ state_dict['model'] = (model[0].sharded_state_dict()
+ if use_dist_ckpt else
+ model[0].state_dict_for_save_checkpoint())
+ else:
+ for i in range(len(model)):
+ mpu.set_virtual_pipeline_model_parallel_rank(i)
+ state_dict['model%d' % i] = (
+ model[i].sharded_state_dict()
+ if use_dist_ckpt else
+ model[i].state_dict_for_save_checkpoint())
+ # Optimizer stuff.
+ if not args.no_save_optim:
+ if optimizer is not None and not optimizer.is_stub_optimizer:
+ state_dict['optimizer'] = (optimizer.sharded_state_dict(state_dict, **(optim_sd_kwargs or {}))
+ if use_dist_ckpt else
+ optimizer.state_dict())
+ if opt_param_scheduler is not None:
+ state_dict['opt_param_scheduler'] = \
+ opt_param_scheduler.state_dict()
+
+ # Rerun state
+ state_dict['rerun_state_machine'] = rerun_state
+
+ # RNG states.
+ if not args.no_save_rng:
+ state_dict["rng_state"] = rng_state
+ return state_dict
+
+
+def _transpose_first_dim(t, num_splits, num_splits_first, model):
+ input_shape = t.size()
+ # We use a self_attention module but the values extracted aren't
+ # specific to self attention so should work for cross attention as well
+ while hasattr(model, 'module'):
+ model = model.module
+ attention_module = model.language_model.encoder.layers[0].self_attention
+ hidden_size_per_attention_head = attention_module.hidden_size_per_attention_head
+ num_attention_heads_per_partition = attention_module.num_attention_heads_per_partition
+ if num_splits_first:
+ """[num_splits * np * hn, h]
+ -->(view) [num_splits, np, hn, h]
+ -->(tranpose) [np, num_splits, hn, h]
+ -->(view) [np * num_splits * hn, h] """
+
+ intermediate_shape = \
+ (num_splits, num_attention_heads_per_partition,
+ hidden_size_per_attention_head) + input_shape[1:]
+
+ t = t.view(*intermediate_shape)
+ t = t.transpose(0, 1).contiguous()
+ else:
+ """[np * hn * num_splits, h]
+ -->(view) [np, hn, num_splits, h]
+ -->(tranpose) [np, num_splits, hn, h]
+ -->(view) [np * num_splits * hn, h] """
+
+ intermediate_shape = \
+ (num_attention_heads_per_partition,
+ hidden_size_per_attention_head, num_splits) +\
+ input_shape[1:]
+
+ t = t.view(*intermediate_shape)
+ t = t.transpose(1, 2).contiguous()
+ t = t.view(*input_shape)
+
+ return t
+
+
+def fix_query_key_value_ordering(model, checkpoint_version):
+ """Fix up query/key/value matrix ordering if checkpoint
+ version is smaller than 2.0
+ """
+ if checkpoint_version < 2.0:
+ if isinstance(model, list):
+ assert len(model)==1
+ model = model[0]
+ for name, param in model.named_parameters():
+ if name.endswith(('.query_key_value.weight', '.query_key_value.bias')):
+ if checkpoint_version == 0:
+ fixed_param = _transpose_first_dim(param.data, 3, True, model)
+ elif checkpoint_version == 1.0:
+ fixed_param = _transpose_first_dim(param.data, 3, False, model)
+ else:
+ print_rank_0(f"Invalid checkpoint version {checkpoint_version}.")
+ sys.exit()
+ param.data.copy_(fixed_param)
+ if name.endswith(('.key_value.weight', '.key_value.bias')):
+ if checkpoint_version == 0:
+ fixed_param = _transpose_first_dim(param.data, 2, True, model)
+ elif checkpoint_version == 1.0:
+ fixed_param = _transpose_first_dim(param.data, 2, False, model)
+ else:
+ print_rank_0(f"Invalid checkpoint version {checkpoint_version}.")
+ sys.exit()
+ param.data.copy_(fixed_param)
+ print_rank_0(" successfully fixed query-key-values ordering for"
+ " checkpoint version {}".format(checkpoint_version))
+
+
+def _get_non_persistent_iteration(non_persistent_global_dir, args, checkpointing_context=None):
+ if args.non_persistent_ckpt_type is None:
+ return -1
+ elif args.non_persistent_ckpt_type == "global":
+ tracker_filename = get_checkpoint_tracker_filename(non_persistent_global_dir)
+ if os.path.isfile(tracker_filename):
+ iteration, release = read_metadata(tracker_filename)
+ if release:
+ raise RuntimeError('Non-persistent checkpoint can\'t be a release checkpoint')
+ else:
+ iteration = -1
+ print_rank_0('WARNING: could not find the metadata file {}'.format(tracker_filename))
+ print_rank_0(' will not load any non-persistent checkpoint')
+ return iteration
+ elif args.non_persistent_ckpt_type == "local":
+ return checkpointing_context['local_checkpoint_manager'].find_latest()
+ else:
+ assert False, 'Please use local or global non-persistent checkpoints' \
+ f'(got: {args.non_persistent_ckpt_type})'
+
+
+def _load_non_persistent_base_checkpoint(
+ non_persistent_global_dir,
+ args,
+ rank0,
+ sharded_state_dict,
+ non_persistent_iteration,
+ checkpointing_context=None,
+):
+ """ Load the base state_dict from a non-persistent distributed checkpoint.
+ Depending on the non_persistent_ckpt_type, different logic may be required.
+ """
+ assert args.non_persistent_ckpt_type is not None
+ if args.non_persistent_ckpt_type == "global":
+ if not rank0:
+ print_rank_0(
+ f'Loading from a non-persistent checkpoint (non-persistent iter {non_persistent_iteration})'
+ )
+ return _load_global_dist_base_checkpoint(
+ non_persistent_global_dir, args, rank0, sharded_state_dict, non_persistent_iteration, False,
+ checkpointing_context=checkpointing_context
+ )
+ elif args.non_persistent_ckpt_type == "local":
+ intermediate_state_dict, checkpoint_name = checkpointing_context[
+ 'local_checkpoint_manager'
+ ].load()
+ state_dict = intermediate_state_dict.to_state_dict(
+ sharded_state_dict,
+ algo=args.non_persistent_local_ckpt_algo,
+ parallelization_group=mpu.get_data_parallel_group(with_context_parallel=True)
+ )
+ return state_dict, checkpoint_name, False, CheckpointType.LOCAL
+ else:
+ assert False, 'Please use local or global non-persistent checkpoints' \
+ f'(got: {args.non_persistent_ckpt_type})'
+
+
+def _load_global_dist_base_checkpoint(
+ load_dir, args, rank0, sharded_state_dict, iteration, release, checkpointing_context=None
+):
+ """ Load the base state_dict from the given directory containing the global distributed checkpoint """
+ if rank0:
+ checkpoint_name = find_checkpoint_rank_0(load_dir, iteration, release)
+ state_dict = dist_checkpointing.load_common_state_dict(checkpoint_name)
+ return state_dict, checkpoint_name, release, CheckpointType.GLOBAL
+
+ if sharded_state_dict is None:
+ assert not args.auto_detect_ckpt_format and not args.use_dist_ckpt, (
+ args.auto_detect_ckpt_format,
+ args.use_dist_ckpt,
+ )
+ raise RuntimeError(
+ 'Detected load from a distributed checkpoint, '
+ 'but neither --use-dist-ckpt nor --auto-detect-ckpt-format is set.'
+ )
+
+ checkpoint_name = get_checkpoint_name(load_dir, iteration, release, return_base_dir=True)
+ load_strategy = get_default_load_sharded_strategy(checkpoint_name)
+ # NOTE: `args.ckpt_fully_parallel_load` applies to both persistent and non-persistent checkpoints.
+ if args.ckpt_fully_parallel_load:
+ load_strategy = FullyParallelLoadStrategyWrapper(
+ load_strategy, mpu.get_data_parallel_group(with_context_parallel=True)
+ )
+ if checkpointing_context is not None:
+ checkpointing_context["load_strategy"] = load_strategy
+ state_dict = dist_checkpointing.load(sharded_state_dict, checkpoint_name,
+ load_strategy, strict=args.dist_ckpt_strictness)
+ return state_dict, checkpoint_name, release, CheckpointType.GLOBAL
+
+
+def _load_base_checkpoint(
+ load_dir,
+ args,
+ rank0=False,
+ sharded_state_dict=None,
+ checkpointing_context=None,
+):
+ """ Load the base state_dict from the given directory
+
+ If rank0 is true, just loads rank 0 checkpoint, ignoring arguments.
+ """
+ # Try to load non-persistent checkpoint first
+ non_persistent_global_dir = (
+ args.non_persistent_global_ckpt_dir
+ if args.non_persistent_global_ckpt_dir or load_dir is None
+ else os.path.join(load_dir, _NON_PERSISTENT_CKPT_SUBDIR)
+ )
+ non_persistent_iteration = _get_non_persistent_iteration(
+ non_persistent_global_dir, args, checkpointing_context
+ )
+ iteration, release = -1, False
+ tracker_filename = 'because load directory is not defined'
+ if load_dir is not None:
+ tracker_filename = get_checkpoint_tracker_filename(load_dir)
+ if os.path.isfile(tracker_filename):
+ iteration, release = read_metadata(tracker_filename)
+ if non_persistent_iteration != -1: # there is a non-persistent checkpoint
+ if non_persistent_iteration >= iteration:
+ return _load_non_persistent_base_checkpoint(
+ non_persistent_global_dir,
+ args,
+ rank0,
+ sharded_state_dict,
+ non_persistent_iteration,
+ checkpointing_context,
+ )
+ else:
+ print_rank_0('WARNING: non-persistent checkpoints are older than persistent checkpoint')
+
+ # Otherwise we are dealing with global checkpoints
+ # If no tracker file, return nothing
+ if iteration == -1:
+ if not rank0:
+ print_rank_0('WARNING: could not find the metadata file {}'.format(tracker_filename))
+ print_rank_0(' will not load any checkpoints and will start from random')
+ # Conditionally exit if checkpoint not found.
+ if args.exit_on_missing_checkpoint:
+ print_rank_0(">> '--exit-on-missing-checkpoint' set ... exiting. <<")
+ if torch.distributed.is_initialized():
+ torch.distributed.barrier()
+ sys.exit()
+
+ return None, "", False, None
+
+ # Determine the type of the checkpoint
+ checkpoint_name = get_checkpoint_name(load_dir, iteration, release, return_base_dir=True)
+ is_dist_ckpt = dist_checkpointing.check_is_distributed_checkpoint(checkpoint_name)
+ if not rank0:
+ dist_infix = "distributed " if is_dist_ckpt else ""
+ if release:
+ print_rank_0(f' loading release {dist_infix}checkpoint from {load_dir}')
+ else:
+ print_rank_0(
+ f' loading {dist_infix}checkpoint from {load_dir} at iteration {iteration}'
+ )
+
+ # Handle global distributed checkpoint
+ if is_dist_ckpt:
+ return _load_global_dist_base_checkpoint(
+ load_dir, args, rank0, sharded_state_dict, iteration, release, checkpointing_context=checkpointing_context
+ )
+ # Handle global legacy checkpoint
+ if rank0:
+ checkpoint_name = find_checkpoint_rank_0(load_dir, iteration, release)
+ else:
+ checkpoint_name = get_checkpoint_name(load_dir, iteration, release, return_base_dir=False)
+ try:
+ state_dict = torch.load(checkpoint_name, map_location='cpu', weights_only=False)
+ except ModuleNotFoundError:
+ from megatron.legacy.fp16_deprecated import loss_scaler
+
+ # For backward compatibility.
+ if not rank0:
+ print_rank_0(' > deserializing using the old code structure ...')
+ sys.modules['fp16.loss_scaler'] = sys.modules['megatron.legacy.fp16_deprecated.loss_scaler']
+ sys.modules['megatron.fp16.loss_scaler'] = sys.modules['megatron.legacy.fp16_deprecated.loss_scaler']
+ sys.modules['megatron.model'] = sys.modules['megatron.legacy.model']
+ state_dict = torch.load(checkpoint_name, map_location='cpu', weights_only=False)
+ sys.modules.pop('fp16.loss_scaler', None)
+ sys.modules.pop('megatron.fp16.loss_scaler', None)
+ sys.modules.pop('megatron.model', None)
+ except Exception as e:
+ print('could not load the checkpoint')
+ print(e)
+ sys.exit()
+
+ return state_dict, checkpoint_name, release, CheckpointType.LEGACY
+
+
+def load_args_from_checkpoint(
+ args, load_arg='load', checkpointing_context=None
+):
+ """Set required arguments from the checkpoint specified in the
+ arguments.
+
+ Will overwrite arguments that have a non-None default value, but
+ will leave any arguments that default to None as set.
+
+ Returns the same args NameSpace with the new values added/updated.
+
+ If no checkpoint is specified in args, or if the checkpoint is
+ there but invalid, the arguments will not be modified
+
+ """
+ load_dir = getattr(args, load_arg)
+
+ if load_dir is None:
+ print_rank_0('No load directory specified, using provided arguments.')
+ return args
+
+ state_dict, checkpoint_name, release, ckpt_type = _load_base_checkpoint(
+ load_dir,
+ args,
+ rank0=True,
+ checkpointing_context=checkpointing_context,
+ )
+
+ # Args.
+ if not state_dict:
+ print_rank_0('Checkpoint not found to provide arguments, using provided arguments.')
+ return args
+
+ if 'args' not in state_dict:
+ print_rank_0('Checkpoint provided does not have arguments saved, using provided arguments.')
+ return args
+
+ checkpoint_args = state_dict['args']
+ checkpoint_version = state_dict.get('checkpoint_version', 0)
+ args.iteration = state_dict['iteration']
+
+ # One-off conversion for foundation models
+ if hasattr(checkpoint_args, 'disable_bias_linear'):
+ setattr(checkpoint_args, 'add_bias_linear', not getattr(checkpoint_args, 'disable_bias_linear'))
+
+ def _set_arg(arg_name, old_arg_name=None, force=False):
+ if not force and getattr(args, arg_name, None) is not None:
+ return
+
+ if old_arg_name is not None:
+ checkpoint_value = getattr(checkpoint_args, old_arg_name, None)
+ else:
+ checkpoint_value = getattr(checkpoint_args, arg_name, None)
+
+ if checkpoint_value is not None:
+ print_rank_0(f"Setting {arg_name} to {checkpoint_value} from checkpoint")
+ setattr(args, arg_name, checkpoint_value)
+ else:
+ print_rank_0(f"Checkpoint did not provide arguments {arg_name}")
+
+ # Model args.
+ _set_arg('num_layers')
+ _set_arg('hidden_size')
+ _set_arg('ffn_hidden_size')
+ _set_arg('seq_length')
+ _set_arg('num_attention_heads')
+ _set_arg('num_query_groups', force=True)
+ _set_arg('group_query_attention', force=True)
+ _set_arg('kv_channels')
+ _set_arg('max_position_embeddings')
+ _set_arg('position_embedding_type', force=True)
+ _set_arg('add_position_embedding', force=True)
+ _set_arg('use_rotary_position_embeddings', force=True)
+ _set_arg('rotary_base', force=True)
+ _set_arg('rotary_percent', force=True)
+ _set_arg('rotary_interleaved', force=True)
+ _set_arg('add_bias_linear', force=True)
+ _set_arg('add_qkv_bias', force=True)
+ _set_arg('squared_relu', force=True)
+ _set_arg('swiglu', force=True)
+ _set_arg('untie_embeddings_and_output_weights', force=True)
+ _set_arg('apply_layernorm_1p', force=True)
+ _set_arg('normalization', force=True)
+ _set_arg('apply_query_key_layer_scaling', force=True)
+ _set_arg('attention_dropout', force=True)
+ _set_arg('hidden_dropout', force=True)
+
+ _set_arg('hybrid_override_pattern', force=True)
+ _set_arg('spec', force=True)
+ _set_arg('hybrid_attention_ratio', force=True)
+ _set_arg('hybrid_mlp_ratio', force=True)
+
+ _set_arg('num_experts', force=True)
+ _set_arg('moe_layer_freq', force=True)
+ _set_arg('moe_ffn_hidden_size', force=True)
+ _set_arg('moe_router_topk', force=True)
+ _set_arg('moe_token_dispatcher_type', force=True)
+ _set_arg('moe_router_pre_softmax', force=True)
+ _set_arg('moe_grouped_gemm', force=True)
+ _set_arg('moe_shared_expert_intermediate_size', force=True)
+
+ # Tokenizer args.
+ _set_arg('tokenizer_type', force=True)
+ # Using checkpoint version might not always be safe (e.g., if running on different cluster).
+ if args.use_tokenizer_model_from_checkpoint_args:
+ _set_arg('tokenizer_model', force=True)
+ _set_arg('tiktoken_pattern', force=True)
+ _set_arg('padded_vocab_size')
+
+ # Checkpoint args.
+ _set_arg('ckpt_format')
+
+ # Model parallelism args.
+ if args.use_mp_args_from_checkpoint_args:
+ if checkpoint_version < 3.0:
+ _set_arg('tensor_model_parallel_size', 'model_parallel_size')
+ else:
+ _set_arg('tensor_model_parallel_size', force=True)
+ _set_arg('pipeline_model_parallel_size', force=True)
+ _set_arg('virtual_pipeline_model_parallel_size', force=True)
+ _set_arg('num_layers_per_virtual_pipeline_stage')
+ _set_arg('expert_model_parallel_size', force=True)
+
+ return args, checkpoint_args
+
+
+def fix_fp8_params_lose_precision_when_loading_dist_ckpt(state_dict):
+ """
+ When "--fp8-param-gather" and "--use-dist-ckpt" are both enabled, the state dict read from
+ dist-checkpoint loses precision (the weights read from checkpoint go through the process of
+ bf16/fp16 -> fp8 -> bf16/fp16). This function is implemented to solve this problem.
+ When "--fp8-param-gather" is disabled, this function doesn't modify anything.
+ """
+ for key in state_dict.keys():
+ if key.startswith('model'):
+ for _, sharded_tensor in state_dict[key].items():
+ if is_float8tensor(sharded_tensor.data):
+ sharded_tensor.data = dequantize_fp8_tensor(sharded_tensor.data).cpu()
+
+
+def load_checkpoint(model, optimizer, opt_param_scheduler, load_arg='load', strict=True,
+ checkpointing_context=None, skip_load_to_model_and_opt=False):
+ """Load a model checkpoint and return the iteration.
+ strict (bool): whether to strictly enforce that the keys in
+ :attr:`state_dict` of the checkpoint match the names of
+ parameters and buffers in model.
+ skip_load_to_model_and_opt (bool): whether to call `load_state_dict`
+ for :attr:`model` and :attr:`optimizer`. In case of running FSDP2
+ or other torch features that uses DTensor in state dict, the tensors
+ are already loaded in-place by `_load_base_checkpoint`.
+ """
+ args = get_args()
+ load_dir = getattr(args, load_arg)
+
+ # Finetuning directories
+ pretrained_dir = getattr(args, 'pretrained_checkpoint', None)
+ if pretrained_dir is not None and not checkpoint_exists(load_dir):
+ print_rank_0(f'Checkpoint file not found in load directory {load_dir} attempting to finetune '
+ f'with checkpoint in {pretrained_dir}')
+ load_dir = pretrained_dir
+ if not checkpoint_exists(load_dir):
+ raise FileNotFoundError("No checkpoint found in load directory or pretrained directory")
+ args.finetune = True
+
+ ddp_model = model
+ print("ddp",ddp_model)
+ model = unwrap_model(ddp_model)
+
+ load_kwargs = {}
+ is_dist_ckpt = False
+ if (
+ args.auto_detect_ckpt_format
+ or args.use_dist_ckpt
+ or args.non_persistent_save_interval is not None
+ ):
+ state_dict, checkpoint_name, release, ckpt_type = _load_base_checkpoint(
+ load_dir,
+ args,
+ rank0=True,
+ checkpointing_context=checkpointing_context,
+ )
+
+ is_dist_ckpt = (
+ ckpt_type == CheckpointType.LOCAL
+ or dist_checkpointing.check_is_distributed_checkpoint(checkpoint_name)
+ )
+ if is_dist_ckpt:
+ ckpt_tp_pp = (
+ state_dict['args'].tensor_model_parallel_size,
+ state_dict['args'].pipeline_model_parallel_size,
+ getattr(state_dict['args'], 'encoder_tensor_model_parallel_size', 0),
+ getattr(state_dict['args'], 'encoder_pipeline_model_parallel_size', 0),
+ )
+ run_tp_pp = (
+ args.tensor_model_parallel_size,
+ args.pipeline_model_parallel_size,
+ # TODO: change this to args.encoder_tensor_model_parallel_size after 30th Nov 24
+ getattr(args, 'encoder_tensor_model_parallel_size', 0),
+ getattr(args, 'encoder_pipeline_model_parallel_size', 0),
+ )
+ mismatch_msg = "(TP, PP, encoder TP, encoder PP) mismatch after resume ({} vs {} from checkpoint)".format(
+ run_tp_pp, ckpt_tp_pp
+ )
+
+ # Determine if RNG state will be loaded
+ if (ckpt_tp_pp == run_tp_pp and not release and not args.finetune and not args.no_load_rng
+ and not getattr(state_dict['args'], 'no_save_rng', False)):
+ gen_sd_rng_state = get_rng_state(True) # we can load the rng state
+ else:
+ gen_sd_rng_state = None
+ if ckpt_tp_pp != run_tp_pp:
+ print_rank_0("{}: RNG state will be ignored".format(mismatch_msg))
+
+ optim_sd_kwargs = dict(is_loading=True)
+ # Determine if optimizer state will be loaded
+ if (not release and not args.finetune and not args.no_load_optim
+ and not getattr(state_dict['args'], 'no_save_optim', False)):
+ gen_sd_optim = optimizer
+ gen_sd_opt_param_scheduler = opt_param_scheduler
+
+ if args.use_distributed_optimizer:
+ optim_sd_kwargs['sharding_type'] = ('fully_sharded_model_space'
+ if getattr(state_dict['args'],
+ 'ckpt_fully_parallel_save', False)
+ else 'dp_zero_gather_scatter')
+ # This is for backwards-compatibility.
+ # Can be removed once 'fully_sharded_bucket_space' loading is removed
+ for maybe_dist_opt_optim_state in (state_dict['optimizer'], *state_dict['optimizer'].values()):
+ if 'param_state_sharding_type' in maybe_dist_opt_optim_state:
+ if maybe_dist_opt_optim_state['param_state_sharding_type'] == 'fully_sharded_bucket_space':
+ print_rank_0('Detected deprecated `fully_sharded_bucket_space` '
+ 'DistributedOptimizer checkpoint format')
+ optim_sd_kwargs['sharding_type'] = \
+ maybe_dist_opt_optim_state['param_state_sharding_type']
+ break
+
+ if ckpt_tp_pp != run_tp_pp and optim_sd_kwargs['sharding_type'] != 'fully_sharded_model_space':
+ raise RuntimeError(f"{mismatch_msg}: not supported for DistributedOptimizer with "
+ f"sharding type {optim_sd_kwargs['sharding_type']}."
+ f" Please use `--ckpt-fully-parallel-save` flag during checkpoint saving.")
+ else:
+ gen_sd_optim = None
+ gen_sd_opt_param_scheduler = None
+
+ # Determine if rerun state will be loaded
+ if (
+ ckpt_tp_pp == run_tp_pp
+ and not release
+ and not args.finetune
+ and 'rerun_state_machine' in state_dict
+ ):
+ rerun_state_machine = get_rerun_state_machine()
+ gen_sd_rerun_state = rerun_state_machine.state_dict(
+ data_iterator=None, use_dist_ckpt=True
+ )
+ else:
+ gen_sd_rerun_state = None
+ if ckpt_tp_pp != run_tp_pp:
+ print_rank_0("{}: Rerun state will be ignored".format(mismatch_msg))
+
+ # [ModelOpt]: IMPORTANT! Restoring modelopt_state (sharded or not) must be performed
+ # after the model instance has been created and before _load_base_checkpoint is called.
+ if has_nvidia_modelopt:
+ if ckpt_type == CheckpointType.LOCAL:
+ print_rank_0('WARNING: Local checkpointing does not support nvidia_modelopt.')
+ elif ckpt_type == CheckpointType.GLOBAL:
+ restore_modelopt_state(model, state_dict)
+ else:
+ restore_sharded_modelopt_state(model, checkpoint_name)
+
+ # [ModelOpt]: Initial loading from non-resume sharded checkpoint to a Distillation Model
+ # will result in key mismatch with loss modules potentially containing parameters, since
+ # it requires generating a state_dict before loading. Here we hide those modules if present.
+ with contextlib.ExitStack() as stack: # Allows multiple context managers for each model shard
+ if args.finetune and hasattr(model[0], "hide_loss_modules"):
+ for m in model:
+ stack.enter_context(m.hide_loss_modules())
+ load_kwargs['sharded_state_dict'] = generate_state_dict(
+ args, model, gen_sd_optim, gen_sd_opt_param_scheduler, gen_sd_rng_state,
+ use_dist_ckpt=True, optim_sd_kwargs=optim_sd_kwargs, rerun_state=gen_sd_rerun_state
+ )
+
+ # When "--fp8-param-gather" is disabled, this function doesn't modify anything.
+ fix_fp8_params_lose_precision_when_loading_dist_ckpt(load_kwargs['sharded_state_dict'])
+
+ state_dict, checkpoint_name, release, ckpt_type = _load_base_checkpoint(
+ load_dir, args, rank0=False, checkpointing_context=checkpointing_context,
+ **load_kwargs
+ )
+
+ # Checkpoint not loaded.
+ if state_dict is None:
+ # Iteration and num_floating_point_operations_so_far default to 0.
+ return 0, 0
+
+ # Set checkpoint version.
+ set_checkpoint_version(state_dict.get('checkpoint_version', 0))
+
+ # Set iteration.
+ if args.finetune or release:
+ iteration = 0
+ else:
+ try:
+ iteration = state_dict['iteration']
+ except KeyError:
+ try: # Backward compatible with older checkpoints
+ iteration = state_dict['total_iters']
+ except KeyError:
+ print_rank_0('A metadata file exists but unable to load '
+ 'iteration from checkpoint {}, exiting'.format(checkpoint_name))
+ sys.exit()
+ num_floating_point_operations_so_far = state_dict.get('num_floating_point_operations_so_far', 0)
+
+ # Check arguments.
+ if load_arg == 'load':
+ assert args.consumed_train_samples == 0
+ assert args.skipped_train_samples == 0
+ assert args.consumed_valid_samples == 0
+ if 'args' in state_dict and not args.finetune:
+ checkpoint_args = state_dict['args']
+ check_checkpoint_args(checkpoint_args)
+ args.consumed_train_samples = getattr(checkpoint_args, 'consumed_train_samples', 0)
+ args.skipped_train_samples = getattr(checkpoint_args, 'skipped_train_samples', 0)
+ update_num_microbatches(consumed_samples=args.consumed_train_samples, verbose=True)
+ args.consumed_valid_samples = getattr(checkpoint_args, 'consumed_valid_samples', 0)
+ else:
+ print_rank_0('could not find arguments in the checkpoint ...')
+
+ # Model.
+ strict = False if args.retro_add_retriever else strict
+ if not skip_load_to_model_and_opt:
+ if len(ddp_model) == 1:
+ ddp_model[0].load_state_dict(state_dict['model'], strict=strict)
+ else:
+ for i in range(len(ddp_model)):
+ mpu.set_virtual_pipeline_model_parallel_rank(i)
+ ddp_model[i].load_state_dict(state_dict['model%d' % i], strict=strict)
+
+ # Fix up query/key/value matrix ordering if needed.
+ checkpoint_version = get_checkpoint_version()
+ print_rank_0(f' checkpoint version {checkpoint_version}')
+ fix_query_key_value_ordering(model, checkpoint_version)
+
+ # Optimizer.
+ if not release and not args.finetune and not args.no_load_optim:
+ try:
+ # Load state dict.
+ if not skip_load_to_model_and_opt and optimizer is not None and not optimizer.is_stub_optimizer:
+ optimizer.load_state_dict(state_dict['optimizer'])
+
+ # Load distributed optimizer's custom parameter state.
+ # For distributed checkpoint it's already loaded in load_state_dict above
+ if args.use_distributed_optimizer and not is_dist_ckpt:
+ # NOTE: this is a manual read of the tracker file.
+ # This code should not be reached when reading from a non_persistent checkpoint
+ assert not is_dist_ckpt
+ tracker_filename = get_checkpoint_tracker_filename(load_dir)
+ iteration, release = read_metadata(tracker_filename)
+ model_checkpoint_name = get_checkpoint_name(load_dir, iteration, release)
+ optim_checkpoint_name = get_distributed_optimizer_checkpoint_name(model_checkpoint_name)
+ optimizer.load_parameter_state(optim_checkpoint_name,
+ update_legacy_format=args.ckpt_convert_update_legacy_dist_opt_format)
+
+ # Load scheduler.
+ if opt_param_scheduler is not None:
+ if 'lr_scheduler' in state_dict: # backward compatbility
+ opt_param_scheduler.load_state_dict(state_dict['lr_scheduler'])
+ else:
+ opt_param_scheduler.load_state_dict(state_dict['opt_param_scheduler'])
+ except KeyError as e:
+ print_rank_0('Unable to load optimizer from checkpoint {}. '
+ 'Specify --no-load-optim or --finetune to prevent '
+ 'attempting to load the optimizer state, '
+ 'exiting ...'.format(checkpoint_name))
+ raise e
+ else:
+ if (args.fp16 or args.bf16) and optimizer is not None:
+ optimizer.reload_model_params()
+
+ # rerun state
+ try:
+ if 'rerun_state_machine' in state_dict:
+ get_rerun_state_machine().load_state_dict(state_dict['rerun_state_machine'])
+ except Exception as e:
+ print(f"Unable to restore RerunMachine from checkpoint: {e}")
+ sys.exit()
+
+ # rng states.
+ if not release and not args.finetune and not args.no_load_rng:
+ try:
+ if 'rng_state' in state_dict:
+ # access rng_state for data parallel rank
+ if args.data_parallel_random_init:
+ rng_state = state_dict['rng_state'][mpu.get_data_parallel_rank()]
+ else:
+ rng_state = state_dict['rng_state'][0]
+ random.setstate(rng_state['random_rng_state'])
+ np.random.set_state(rng_state['np_rng_state'])
+ torch.set_rng_state(rng_state['torch_rng_state'])
+ torch.cuda.set_rng_state(rng_state['cuda_rng_state'])
+ # Check for empty states array
+ if not rng_state['rng_tracker_states']:
+ raise KeyError
+ tensor_parallel.get_cuda_rng_tracker().set_states(
+ rng_state['rng_tracker_states'])
+ else: # backward compatability
+ random.setstate(state_dict['random_rng_state'])
+ np.random.set_state(state_dict['np_rng_state'])
+ torch.set_rng_state(state_dict['torch_rng_state'])
+ torch.cuda.set_rng_state(state_dict['cuda_rng_state'])
+ # Check for empty states array
+ if not state_dict['rng_tracker_states']:
+ raise KeyError
+ tensor_parallel.get_cuda_rng_tracker().set_states(
+ state_dict['rng_tracker_states'])
+ except KeyError:
+ print_rank_0('Unable to load rng state from checkpoint {}. '
+ 'Specify --no-load-rng or --finetune to prevent '
+ 'attempting to load the rng state, '
+ 'exiting ...'.format(checkpoint_name))
+ sys.exit()
+
+ # Some utilities want to load a checkpoint without distributed being initialized
+ if torch.distributed.is_initialized():
+ torch.distributed.barrier()
+
+ print_rank_0(f' successfully loaded checkpoint from {load_dir} '
+ f'[ t {mpu.get_tensor_model_parallel_rank() + 1}/{mpu.get_tensor_model_parallel_world_size()}, '
+ f'p {mpu.get_pipeline_model_parallel_rank() + 1}/{mpu.get_pipeline_model_parallel_world_size()} ] '
+ f'at iteration {iteration}')
+
+ # Additional callback for wandb (last rank)
+ if not torch.distributed.is_initialized() \
+ or is_last_rank():
+ wandb_utils.on_load_checkpoint_success(checkpoint_name, load_dir)
+
+ torch.cuda.empty_cache()
+
+ if iteration > 0:
+ # Notify FT that a checkpoint was loaded.
+ is_local_chkpt = (ckpt_type == CheckpointType.LOCAL)
+ ft_integration.on_checkpoint_loaded(is_local_chkpt=is_local_chkpt)
+
+ return iteration, num_floating_point_operations_so_far
+
+
+def load_biencoder_checkpoint(model, only_query_model=False,
+ only_context_model=False, custom_load_path=None):
+ """
+ selectively load retrieval models for indexing/retrieving
+ from saved checkpoints
+ """
+
+ args = get_args()
+
+ model = unwrap_model(model)
+
+ load_path = custom_load_path if custom_load_path is not None else args.load
+
+ tracker_filename = get_checkpoint_tracker_filename(load_path)
+ with open(tracker_filename, 'r') as f:
+ iteration = int(f.read().strip())
+
+ checkpoint_name = get_checkpoint_name(load_path, iteration,
+ args.use_distributed_optimizer,
+ release=False)
+
+ if mpu.get_data_parallel_rank() == 0:
+ print('global rank {} is loading checkpoint {}'.format(
+ torch.distributed.get_rank(), checkpoint_name))
+
+ state_dict = torch.load(checkpoint_name, map_location='cpu', weights_only=False)
+ ret_state_dict = state_dict['model']
+
+ if only_query_model:
+ ret_state_dict.pop('context_model')
+ if only_context_model:
+ ret_state_dict.pop('query_model')
+
+ assert len(model) == 1
+ model[0].load_state_dict(ret_state_dict)
+ torch.distributed.barrier()
+
+ if mpu.get_data_parallel_rank() == 0:
+ print(' successfully loaded {}'.format(checkpoint_name))
+
+ return model
diff --git a/aiak_megatron/megatron/training/training.py b/aiak_megatron/megatron/training/training.py
index 0a6d4b82..78f1e4d4 100644
--- a/aiak_megatron/megatron/training/training.py
+++ b/aiak_megatron/megatron/training/training.py
@@ -413,6 +413,21 @@ def pretrain(
iteration = 0
if args.do_train and args.train_iters > 0:
+ # Print parameter trainability status
+ print_rank_0('\n' + '=' * 80)
+ print_rank_0("training with the following parameter status:")
+ print_rank_0('=' * 80)
+ model_to_check = model[0] if isinstance(model, list) else model
+ trainable_params = 0
+ frozen_params = 0
+ for name, param in model_to_check.named_parameters():
+ status = "TRAINABLE" if param.requires_grad else "FROZEN"
+ print_rank_0(f"{status:12s} | {name:80s} | shape: {str(tuple(param.shape)):30s} | dtype: {param.dtype}")
+ if param.requires_grad:
+ trainable_params += param.numel()
+ else:
+ frozen_params += param.numel()
+
iteration, num_floating_point_operations_so_far = train(
forward_step_func,
model, optimizer, opt_param_scheduler,
diff --git a/aiak_megatron/megatron/training/utils.py b/aiak_megatron/megatron/training/utils.py
index 61b18609..97199e9e 100644
--- a/aiak_megatron/megatron/training/utils.py
+++ b/aiak_megatron/megatron/training/utils.py
@@ -381,6 +381,9 @@ def print_rank_last(message):
def get_device_arch_version():
"""Returns GPU arch version (8: Ampere, 9: Hopper, 10: Blackwell, ...)"""
+ # On AMD/ROCm, return a high version to skip CUDA-specific checks
+ if torch.version.hip is not None:
+ return 99 # Return high value to skip CUDA_DEVICE_MAX_CONNECTIONS checks
return torch.cuda.get_device_properties(torch.device("cuda:0")).major
diff --git a/aiak_megatron/megatron/training/wandb_utils.py b/aiak_megatron/megatron/training/wandb_utils.py
index ed3a55aa..ff339613 100644
--- a/aiak_megatron/megatron/training/wandb_utils.py
+++ b/aiak_megatron/megatron/training/wandb_utils.py
@@ -34,7 +34,9 @@ def on_save_checkpoint_success(checkpoint_path: str, tracker_filename: str, save
metadata = {"iteration": iteration}
artifact_name, artifact_version = _get_artifact_name_and_version(Path(save_dir), Path(checkpoint_path))
artifact = wandb_writer.Artifact(artifact_name, type="model", metadata=metadata)
- artifact.add_reference(f"file://{checkpoint_path}", checksum=False)
+ # Convert to absolute path for W&B
+ abs_checkpoint_path = Path(checkpoint_path).resolve()
+ artifact.add_reference(f"file://{abs_checkpoint_path}", checksum=False)
artifact.add_file(tracker_filename)
wandb_writer.run.log_artifact(artifact, aliases=[artifact_version])
wandb_tracker_filename = _get_wandb_artifact_tracker_filename(save_dir)
diff --git a/aiak_megatron/megatron_core.egg-info/PKG-INFO b/aiak_megatron/megatron_core.egg-info/PKG-INFO
new file mode 100644
index 00000000..53970e3b
--- /dev/null
+++ b/aiak_megatron/megatron_core.egg-info/PKG-INFO
@@ -0,0 +1,332 @@
+Metadata-Version: 2.4
+Name: megatron-core
+Version: 0.12.0rc0
+Summary: Megatron Core - a library for efficient and scalable training of transformer based models
+Home-page: https://github.com/NVIDIA/Megatron-LM/megatron/core
+Download-URL: https://github.com/NVIDIA/Megatron-LM/releases
+Author: NVIDIA
+Author-email: NVIDIA
+Maintainer: NVIDIA
+Maintainer-email: NVIDIA
+License: The following applies to all files unless otherwise noted:
+
+ # Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+ #
+ # Redistribution and use in source and binary forms, with or without
+ # modification, are permitted provided that the following conditions
+ # are met:
+ # * Redistributions of source code must retain the above copyright
+ # notice, this list of conditions and the following disclaimer.
+ # * Redistributions in binary form must reproduce the above copyright
+ # notice, this list of conditions and the following disclaimer in the
+ # documentation and/or other materials provided with the distribution.
+ # * Neither the name of NVIDIA CORPORATION nor the names of its
+ # contributors may be used to endorse or promote products derived
+ # from this software without specific prior written permission.
+ #
+ # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
+ # EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+ # PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
+ # CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+ # EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+ # PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
+ # PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+ # OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+ --
+
+ This repository also contains code from Hugging Face Inc., Google Research,
+ Facebook (from their Fairseq, Dino, and ParlAI projects), Microsoft (from their
+ Swin-Transformer project), Philip Popien, the Mamba project (Tri Dao and
+ Albert Gu), and the Triton language and compiler project (Philippe Tillet and
+ OpenAI). Files from these organizations have notices at the top of each file.
+ Below are licenses used in those files, as indicated.
+
+
+ --------------------------------------------------------------------------------------
+ -- LICENSE FOR Facebook, huggingface, Google Research, LLaVA, Mamba, and vLLM code --
+
+
+ Apache License
+ Version 2.0, January 2004
+ http://www.apache.org/licenses/
+
+ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+ 1. Definitions.
+
+ "License" shall mean the terms and conditions for use, reproduction,
+ and distribution as defined by Sections 1 through 9 of this document.
+
+ "Licensor" shall mean the copyright owner or entity authorized by
+ the copyright owner that is granting the License.
+
+ "Legal Entity" shall mean the union of the acting entity and all
+ other entities that control, are controlled by, or are under common
+ control with that entity. For the purposes of this definition,
+ "control" means (i) the power, direct or indirect, to cause the
+ direction or management of such entity, whether by contract or
+ otherwise, or (ii) ownership of fifty percent (50%) or more of the
+ outstanding shares, or (iii) beneficial ownership of such entity.
+
+ "You" (or "Your") shall mean an individual or Legal Entity
+ exercising permissions granted by this License.
+
+ "Source" form shall mean the preferred form for making modifications,
+ including but not limited to software source code, documentation
+ source, and configuration files.
+
+ "Object" form shall mean any form resulting from mechanical
+ transformation or translation of a Source form, including but
+ not limited to compiled object code, generated documentation,
+ and conversions to other media types.
+
+ "Work" shall mean the work of authorship, whether in Source or
+ Object form, made available under the License, as indicated by a
+ copyright notice that is included in or attached to the work
+ (an example is provided in the Appendix below).
+
+ "Derivative Works" shall mean any work, whether in Source or Object
+ form, that is based on (or derived from) the Work and for which the
+ editorial revisions, annotations, elaborations, or other modifications
+ represent, as a whole, an original work of authorship. For the purposes
+ of this License, Derivative Works shall not include works that remain
+ separable from, or merely link (or bind by name) to the interfaces of,
+ the Work and Derivative Works thereof.
+
+ "Contribution" shall mean any work of authorship, including
+ the original version of the Work and any modifications or additions
+ to that Work or Derivative Works thereof, that is intentionally
+ submitted to Licensor for inclusion in the Work by the copyright owner
+ or by an individual or Legal Entity authorized to submit on behalf of
+ the copyright owner. For the purposes of this definition, "submitted"
+ means any form of electronic, verbal, or written communication sent
+ to the Licensor or its representatives, including but not limited to
+ communication on electronic mailing lists, source code control systems,
+ and issue tracking systems that are managed by, or on behalf of, the
+ Licensor for the purpose of discussing and improving the Work, but
+ excluding communication that is conspicuously marked or otherwise
+ designated in writing by the copyright owner as "Not a Contribution."
+
+ "Contributor" shall mean Licensor and any individual or Legal Entity
+ on behalf of whom a Contribution has been received by Licensor and
+ subsequently incorporated within the Work.
+
+ 2. Grant of Copyright License. Subject to the terms and conditions of
+ this License, each Contributor hereby grants to You a perpetual,
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+ copyright license to reproduce, prepare Derivative Works of,
+ publicly display, publicly perform, sublicense, and distribute the
+ Work and such Derivative Works in Source or Object form.
+
+ 3. Grant of Patent License. Subject to the terms and conditions of
+ this License, each Contributor hereby grants to You a perpetual,
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+ (except as stated in this section) patent license to make, have made,
+ use, offer to sell, sell, import, and otherwise transfer the Work,
+ where such license applies only to those patent claims licensable
+ by such Contributor that are necessarily infringed by their
+ Contribution(s) alone or by combination of their Contribution(s)
+ with the Work to which such Contribution(s) was submitted. If You
+ institute patent litigation against any entity (including a
+ cross-claim or counterclaim in a lawsuit) alleging that the Work
+ or a Contribution incorporated within the Work constitutes direct
+ or contributory patent infringement, then any patent licenses
+ granted to You under this License for that Work shall terminate
+ as of the date such litigation is filed.
+
+ 4. Redistribution. You may reproduce and distribute copies of the
+ Work or Derivative Works thereof in any medium, with or without
+ modifications, and in Source or Object form, provided that You
+ meet the following conditions:
+
+ (a) You must give any other recipients of the Work or
+ Derivative Works a copy of this License; and
+
+ (b) You must cause any modified files to carry prominent notices
+ stating that You changed the files; and
+
+ (c) You must retain, in the Source form of any Derivative Works
+ that You distribute, all copyright, patent, trademark, and
+ attribution notices from the Source form of the Work,
+ excluding those notices that do not pertain to any part of
+ the Derivative Works; and
+
+ (d) If the Work includes a "NOTICE" text file as part of its
+ distribution, then any Derivative Works that You distribute must
+ include a readable copy of the attribution notices contained
+ within such NOTICE file, excluding those notices that do not
+ pertain to any part of the Derivative Works, in at least one
+ of the following places: within a NOTICE text file distributed
+ as part of the Derivative Works; within the Source form or
+ documentation, if provided along with the Derivative Works; or,
+ within a display generated by the Derivative Works, if and
+ wherever such third-party notices normally appear. The contents
+ of the NOTICE file are for informational purposes only and
+ do not modify the License. You may add Your own attribution
+ notices within Derivative Works that You distribute, alongside
+ or as an addendum to the NOTICE text from the Work, provided
+ that such additional attribution notices cannot be construed
+ as modifying the License.
+
+ You may add Your own copyright statement to Your modifications and
+ may provide additional or different license terms and conditions
+ for use, reproduction, or distribution of Your modifications, or
+ for any such Derivative Works as a whole, provided Your use,
+ reproduction, and distribution of the Work otherwise complies with
+ the conditions stated in this License.
+
+ 5. Submission of Contributions. Unless You explicitly state otherwise,
+ any Contribution intentionally submitted for inclusion in the Work
+ by You to the Licensor shall be under the terms and conditions of
+ this License, without any additional terms or conditions.
+ Notwithstanding the above, nothing herein shall supersede or modify
+ the terms of any separate license agreement you may have executed
+ with Licensor regarding such Contributions.
+
+ 6. Trademarks. This License does not grant permission to use the trade
+ names, trademarks, service marks, or product names of the Licensor,
+ except as required for reasonable and customary use in describing the
+ origin of the Work and reproducing the content of the NOTICE file.
+
+ 7. Disclaimer of Warranty. Unless required by applicable law or
+ agreed to in writing, Licensor provides the Work (and each
+ Contributor provides its Contributions) on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+ implied, including, without limitation, any warranties or conditions
+ of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+ PARTICULAR PURPOSE. You are solely responsible for determining the
+ appropriateness of using or redistributing the Work and assume any
+ risks associated with Your exercise of permissions under this License.
+
+ 8. Limitation of Liability. In no event and under no legal theory,
+ whether in tort (including negligence), contract, or otherwise,
+ unless required by applicable law (such as deliberate and grossly
+ negligent acts) or agreed to in writing, shall any Contributor be
+ liable to You for damages, including any direct, indirect, special,
+ incidental, or consequential damages of any character arising as a
+ result of this License or out of the use or inability to use the
+ Work (including but not limited to damages for loss of goodwill,
+ work stoppage, computer failure or malfunction, or any and all
+ other commercial damages or losses), even if such Contributor
+ has been advised of the possibility of such damages.
+
+ 9. Accepting Warranty or Additional Liability. While redistributing
+ the Work or Derivative Works thereof, You may choose to offer,
+ and charge a fee for, acceptance of support, warranty, indemnity,
+ or other liability obligations and/or rights consistent with this
+ License. However, in accepting such obligations, You may act only
+ on Your own behalf and on Your sole responsibility, not on behalf
+ of any other Contributor, and only if You agree to indemnify,
+ defend, and hold each Contributor harmless for any liability
+ incurred by, or claims asserted against, such Contributor by reason
+ of your accepting any such warranty or additional liability.
+
+ END OF TERMS AND CONDITIONS
+
+ APPENDIX: How to apply the Apache License to your work.
+
+ To apply the Apache License to your work, attach the following
+ boilerplate notice, with the fields enclosed by brackets "[]"
+ replaced with your own identifying information. (Don't include
+ the brackets!) The text should be enclosed in the appropriate
+ comment syntax for the file format. We also recommend that a
+ file or class name and description of purpose be included on the
+ same "printed page" as the copyright notice for easier
+ identification within third-party archives.
+
+ Copyright [yyyy] [name of copyright owner]
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+
+ --------------------------------------------------------------------------------
+ LICENSE FOR
+ Facebook, Inc. and its affiliates,
+ Meta Platforms, Inc. and its affiliates,
+ Microsoft Corporation,
+ OpenGVLab/InternVL,
+ Triton language and compiler,
+ and DeepSeek.
+
+ MIT License
+
+ Permission is hereby granted, free of charge, to any person obtaining a copy
+ of this software and associated documentation files (the "Software"), to deal
+ in the Software without restriction, including without limitation the rights
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ copies of the Software, and to permit persons to whom the Software is
+ furnished to do so, subject to the following conditions:
+
+ The above copyright notice and this permission notice shall be included in all
+ copies or substantial portions of the Software.
+
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ SOFTWARE.
+
+
+Project-URL: Download, https://github.com/NVIDIA/Megatron-LM/releases
+Project-URL: Homepage, https://github.com/NVIDIA/Megatron-LM/megatron/core
+Keywords: NLP,NLU,deep,gpu,language,learning,learning,machine,nvidia,pytorch,torch,transformer
+Classifier: Development Status :: 5 - Production/Stable
+Classifier: Environment :: Console
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Information Technology
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: BSD License
+Classifier: Natural Language :: English
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Scientific/Engineering :: Image Recognition
+Classifier: Topic :: Scientific/Engineering :: Mathematics
+Classifier: Topic :: Scientific/Engineering
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Classifier: Topic :: Software Development :: Libraries
+Classifier: Topic :: Utilities
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: einops
+Requires-Dist: flask-restful
+Requires-Dist: nltk
+Requires-Dist: pytest
+Requires-Dist: pytest_asyncio
+Requires-Dist: pytest-cov
+Requires-Dist: pytest_mock
+Requires-Dist: pytest-random-order
+Requires-Dist: sentencepiece
+Requires-Dist: tiktoken
+Requires-Dist: wrapt
+Requires-Dist: zarr
+Requires-Dist: wandb
+Requires-Dist: tensorstore!=0.1.46,!=0.1.72
+Requires-Dist: torch
+Requires-Dist: nvidia-modelopt[torch]>=0.19.0; sys_platform != "darwin"
+Requires-Dist: nvidia-resiliency-ext; sys_platform != "darwin"
+Requires-Dist: torch
+Requires-Dist: packaging
+Dynamic: author
+Dynamic: download-url
+Dynamic: home-page
+Dynamic: license-file
+Dynamic: maintainer
+Dynamic: requires-dist
diff --git a/aiak_megatron/megatron_core.egg-info/SOURCES.txt b/aiak_megatron/megatron_core.egg-info/SOURCES.txt
new file mode 100644
index 00000000..d13b8e94
--- /dev/null
+++ b/aiak_megatron/megatron_core.egg-info/SOURCES.txt
@@ -0,0 +1,305 @@
+LICENSE
+MANIFEST.in
+pyproject.toml
+setup.py
+megatron/core/README.md
+megatron/core/__init__.py
+megatron/core/config.py
+megatron/core/config_logger.py
+megatron/core/enums.py
+megatron/core/fp8_utils.py
+megatron/core/inference_params.py
+megatron/core/jit.py
+megatron/core/model_parallel_config.py
+megatron/core/num_microbatches_calculator.py
+megatron/core/optimizer_param_scheduler.py
+megatron/core/package_info.py
+megatron/core/packed_seq_params.py
+megatron/core/parallel_state.py
+megatron/core/requirements.txt
+megatron/core/rerun_state_machine.py
+megatron/core/timers.py
+megatron/core/utils.py
+megatron/core/datasets/__init__.py
+megatron/core/datasets/bert_dataset.py
+megatron/core/datasets/blended_dataset.py
+megatron/core/datasets/blended_megatron_dataset_builder.py
+megatron/core/datasets/blended_megatron_dataset_config.py
+megatron/core/datasets/gpt_dataset.py
+megatron/core/datasets/helpers.cpp
+megatron/core/datasets/helpers.py
+megatron/core/datasets/indexed_dataset.py
+megatron/core/datasets/masked_dataset.py
+megatron/core/datasets/megatron_dataset.py
+megatron/core/datasets/megatron_tokenizer.py
+megatron/core/datasets/multimodal_dataset.py
+megatron/core/datasets/t5_dataset.py
+megatron/core/datasets/utils.py
+megatron/core/datasets/utils_s3.py
+megatron/core/datasets/retro/__init__.py
+megatron/core/datasets/retro/external_libs.py
+megatron/core/datasets/retro/utils.py
+megatron/core/datasets/retro/config/__init__.py
+megatron/core/datasets/retro/config/bert_embedders.py
+megatron/core/datasets/retro/config/config.py
+megatron/core/datasets/retro/config/gpt_chunk_datasets.py
+megatron/core/datasets/retro/config/tokenizers.py
+megatron/core/datasets/retro/db/__init__.py
+megatron/core/datasets/retro/db/build.py
+megatron/core/datasets/retro/db/dataset.py
+megatron/core/datasets/retro/db/utils.py
+megatron/core/datasets/retro/index/__init__.py
+megatron/core/datasets/retro/index/build.py
+megatron/core/datasets/retro/index/factory.py
+megatron/core/datasets/retro/index/index.py
+megatron/core/datasets/retro/index/utils.py
+megatron/core/datasets/retro/index/validate.py
+megatron/core/datasets/retro/index/indexes/__init__.py
+megatron/core/datasets/retro/index/indexes/faiss_base.py
+megatron/core/datasets/retro/index/indexes/faiss_par_add.py
+megatron/core/datasets/retro/query/__init__.py
+megatron/core/datasets/retro/query/gpt_chunk_dataset.py
+megatron/core/datasets/retro/query/multi_split_gpt_dataset.py
+megatron/core/datasets/retro/query/query.py
+megatron/core/datasets/retro/query/retro_dataset.py
+megatron/core/datasets/retro/query/utils.py
+megatron/core/dist_checkpointing/__init__.py
+megatron/core/dist_checkpointing/core.py
+megatron/core/dist_checkpointing/dict_utils.py
+megatron/core/dist_checkpointing/exchange_utils.py
+megatron/core/dist_checkpointing/mapping.py
+megatron/core/dist_checkpointing/optimizer.py
+megatron/core/dist_checkpointing/serialization.py
+megatron/core/dist_checkpointing/state_dict_utils.py
+megatron/core/dist_checkpointing/tensor_aware_state_dict.py
+megatron/core/dist_checkpointing/utils.py
+megatron/core/dist_checkpointing/validation.py
+megatron/core/dist_checkpointing/strategies/__init__.py
+megatron/core/dist_checkpointing/strategies/async_utils.py
+megatron/core/dist_checkpointing/strategies/base.py
+megatron/core/dist_checkpointing/strategies/cached_metadata_filesystem_reader.py
+megatron/core/dist_checkpointing/strategies/common.py
+megatron/core/dist_checkpointing/strategies/filesystem_async.py
+megatron/core/dist_checkpointing/strategies/fully_parallel.py
+megatron/core/dist_checkpointing/strategies/resharding.py
+megatron/core/dist_checkpointing/strategies/state_dict_saver.py
+megatron/core/dist_checkpointing/strategies/tensorstore.py
+megatron/core/dist_checkpointing/strategies/torch.py
+megatron/core/dist_checkpointing/strategies/two_stage.py
+megatron/core/dist_checkpointing/strategies/zarr.py
+megatron/core/distributed/__init__.py
+megatron/core/distributed/data_parallel_base.py
+megatron/core/distributed/distributed_data_parallel.py
+megatron/core/distributed/distributed_data_parallel_config.py
+megatron/core/distributed/finalize_model_grads.py
+megatron/core/distributed/param_and_grad_buffer.py
+megatron/core/distributed/torch_fully_sharded_data_parallel.py
+megatron/core/distributed/custom_fsdp/__init__.py
+megatron/core/distributed/custom_fsdp/fully_sharded_data_parallel.py
+megatron/core/distributed/custom_fsdp/param_and_grad_buffer.py
+megatron/core/export/__init__.py
+megatron/core/export/data_type.py
+megatron/core/export/export_config.py
+megatron/core/export/model_type.py
+megatron/core/export/trtllm/__init__.py
+megatron/core/export/trtllm/trt_model_config.py
+megatron/core/export/trtllm/trt_model_type.py
+megatron/core/export/trtllm/trtllm_helper.py
+megatron/core/export/trtllm/trtllm_layers.py
+megatron/core/export/trtllm/engine_builder/__init__.py
+megatron/core/export/trtllm/engine_builder/trtllm_engine_builder.py
+megatron/core/export/trtllm/model_to_trllm_mapping/__init__.py
+megatron/core/export/trtllm/model_to_trllm_mapping/default_conversion_dict.py
+megatron/core/export/trtllm/trtllm_weights_converter/__init__.py
+megatron/core/export/trtllm/trtllm_weights_converter/distributed_trtllm_model_weights_converter.py
+megatron/core/export/trtllm/trtllm_weights_converter/single_device_trtllm_model_weights_converter.py
+megatron/core/extensions/__init__.py
+megatron/core/extensions/transformer_engine.py
+megatron/core/extensions/transformer_engine2.py
+megatron/core/extensions/transformer_engine_org.py
+megatron/core/fusions/__init__.py
+megatron/core/fusions/fused_bias_dropout.py
+megatron/core/fusions/fused_bias_geglu.py
+megatron/core/fusions/fused_bias_gelu.py
+megatron/core/fusions/fused_bias_swiglu.py
+megatron/core/fusions/fused_cross_entropy.py
+megatron/core/fusions/fused_indices_converter.py
+megatron/core/fusions/fused_layer_norm.py
+megatron/core/fusions/fused_layer_norm_org.py
+megatron/core/fusions/fused_pad_routing_map.py
+megatron/core/fusions/fused_softmax.py
+megatron/core/inference/__init__.py
+megatron/core/inference/async_stream.py
+megatron/core/inference/common_inference_params.py
+megatron/core/inference/communication_utils.py
+megatron/core/inference/inference_request.py
+megatron/core/inference/sampling_params.py
+megatron/core/inference/scheduler.py
+megatron/core/inference/utils.py
+megatron/core/inference/ammo_support/__init__.py
+megatron/core/inference/ammo_support/gpt/__init__.py
+megatron/core/inference/ammo_support/gpt/model_specs.py
+megatron/core/inference/ammo_support/gpt/state_dict_hooks.py
+megatron/core/inference/engines/__init__.py
+megatron/core/inference/engines/abstract_engine.py
+megatron/core/inference/engines/mcore_engine.py
+megatron/core/inference/gpt/__init__.py
+megatron/core/inference/gpt/model_specs.py
+megatron/core/inference/gpt/state_dict_hooks.py
+megatron/core/inference/model_inference_wrappers/__init__.py
+megatron/core/inference/model_inference_wrappers/abstract_model_inference_wrapper.py
+megatron/core/inference/model_inference_wrappers/inference_wrapper_config.py
+megatron/core/inference/model_inference_wrappers/gpt/__init__.py
+megatron/core/inference/model_inference_wrappers/gpt/gpt_inference_wrapper.py
+megatron/core/inference/model_inference_wrappers/multimodal/vlm_inference_wrapper.py
+megatron/core/inference/model_inference_wrappers/t5/__init__.py
+megatron/core/inference/model_inference_wrappers/t5/t5_inference_wrapper.py
+megatron/core/inference/modelopt_support/__init__.py
+megatron/core/inference/modelopt_support/gpt/__init__.py
+megatron/core/inference/modelopt_support/gpt/model_specs.py
+megatron/core/inference/modelopt_support/gpt/state_dict_hooks.py
+megatron/core/inference/modelopt_support/mamba/__init__.py
+megatron/core/inference/modelopt_support/mamba/model_specs.py
+megatron/core/inference/text_generation_controllers/__init__.py
+megatron/core/inference/text_generation_controllers/encoder_decoder_text_generation_controller.py
+megatron/core/inference/text_generation_controllers/simple_text_generation_controller.py
+megatron/core/inference/text_generation_controllers/text_generation_controller.py
+megatron/core/inference/text_generation_controllers/vlm_text_generation_controller.py
+megatron/core/models/__init__.py
+megatron/core/models/T5/__init__.py
+megatron/core/models/T5/t5_model.py
+megatron/core/models/T5/t5_spec.py
+megatron/core/models/bert/__init__.py
+megatron/core/models/bert/bert_layer_specs.py
+megatron/core/models/bert/bert_lm_head.py
+megatron/core/models/bert/bert_model.py
+megatron/core/models/bert/pooler.py
+megatron/core/models/common/__init__.py
+megatron/core/models/common/embeddings/__init__.py
+megatron/core/models/common/embeddings/language_model_embedding.py
+megatron/core/models/common/embeddings/relative_pos_embedding.py
+megatron/core/models/common/embeddings/rope_utils.py
+megatron/core/models/common/embeddings/rotary_pos_embedding.py
+megatron/core/models/common/embeddings/yarn_rotary_pos_embedding.py
+megatron/core/models/common/language_module/__init__.py
+megatron/core/models/common/language_module/language_module.py
+megatron/core/models/common/vision_module/__init__.py
+megatron/core/models/common/vision_module/vision_module.py
+megatron/core/models/gpt/__init__.py
+megatron/core/models/gpt/gpt_layer_specs.py
+megatron/core/models/gpt/gpt_model.py
+megatron/core/models/gpt/moe_module_specs.py
+megatron/core/models/huggingface/__init__.py
+megatron/core/models/huggingface/clip_model.py
+megatron/core/models/huggingface/module.py
+megatron/core/models/huggingface/qwen_model.py
+megatron/core/models/mamba/__init__.py
+megatron/core/models/mamba/mamba_layer_specs.py
+megatron/core/models/mamba/mamba_model.py
+megatron/core/models/multimodal/__init__.py
+megatron/core/models/multimodal/context_parallel.py
+megatron/core/models/multimodal/llava_model.py
+megatron/core/models/multimodal/llava_spec.py
+megatron/core/models/retro/__init__.py
+megatron/core/models/retro/base_attention.py
+megatron/core/models/retro/config.py
+megatron/core/models/retro/decoder_attention.py
+megatron/core/models/retro/decoder_spec.py
+megatron/core/models/retro/encoder_attention.py
+megatron/core/models/retro/encoder_spec.py
+megatron/core/models/retro/model.py
+megatron/core/models/retro/utils.py
+megatron/core/models/vision/__init__.py
+megatron/core/models/vision/clip_vit_model.py
+megatron/core/models/vision/multimodal_projector.py
+megatron/core/models/vision/radio.py
+megatron/core/models/vision/vit_layer_specs.py
+megatron/core/optimizer/__init__.py
+megatron/core/optimizer/clip_grads.py
+megatron/core/optimizer/distrib_optimizer.py
+megatron/core/optimizer/grad_scaler.py
+megatron/core/optimizer/optimizer.py
+megatron/core/optimizer/optimizer_config.py
+megatron/core/optimizer/cpu_offloading/__init__.py
+megatron/core/optimizer/cpu_offloading/hybrid_optimizer.py
+megatron/core/pipeline_parallel/__init__.py
+megatron/core/pipeline_parallel/combined_1f1b.py
+megatron/core/pipeline_parallel/p2p_communication.py
+megatron/core/pipeline_parallel/schedules.py
+megatron/core/post_training/__init__.py
+megatron/core/post_training/modelopt/__init__.py
+megatron/core/post_training/modelopt/layers.py
+megatron/core/post_training/modelopt/gpt/__init__.py
+megatron/core/post_training/modelopt/gpt/model_specs.py
+megatron/core/post_training/modelopt/gpt/state_dict_hooks.py
+megatron/core/post_training/modelopt/mamba/__init__.py
+megatron/core/post_training/modelopt/mamba/model_specs.py
+megatron/core/ssm/__init__.py
+megatron/core/ssm/mamba_block.py
+megatron/core/ssm/mamba_hybrid_layer_allocation.py
+megatron/core/ssm/mamba_layer.py
+megatron/core/ssm/mamba_mixer.py
+megatron/core/ssm/triton_cache_manager.py
+megatron/core/tensor_parallel/__init__.py
+megatron/core/tensor_parallel/cross_entropy.py
+megatron/core/tensor_parallel/data.py
+megatron/core/tensor_parallel/layers.py
+megatron/core/tensor_parallel/mappings.py
+megatron/core/tensor_parallel/random.py
+megatron/core/tensor_parallel/utils.py
+megatron/core/transformer/__init__.py
+megatron/core/transformer/attention.py
+megatron/core/transformer/cuda_graphs.py
+megatron/core/transformer/dot_product_attention.py
+megatron/core/transformer/dot_product_attention_no_te.py
+megatron/core/transformer/enums.py
+megatron/core/transformer/identity_op.py
+megatron/core/transformer/mlp.py
+megatron/core/transformer/module.py
+megatron/core/transformer/multi_latent_attention.py
+megatron/core/transformer/spec_utils.py
+megatron/core/transformer/torch_layer_norm.py
+megatron/core/transformer/torch_norm.py
+megatron/core/transformer/transformer_block.py
+megatron/core/transformer/transformer_config.py
+megatron/core/transformer/transformer_layer.py
+megatron/core/transformer/utils.py
+megatron/core/transformer/custom_layers/__init__.py
+megatron/core/transformer/custom_layers/transformer_engine.py
+megatron/core/transformer/moe/__init__.py
+megatron/core/transformer/moe/experts.py
+megatron/core/transformer/moe/fused_a2a.py
+megatron/core/transformer/moe/grouped_gemm_util.py
+megatron/core/transformer/moe/legacy_a2a_token_dispatcher.py
+megatron/core/transformer/moe/moe_layer.py
+megatron/core/transformer/moe/moe_utils.py
+megatron/core/transformer/moe/router.py
+megatron/core/transformer/moe/shared_experts.py
+megatron/core/transformer/moe/token_dispatcher.py
+megatron/core/transformer/moe/upcycling_utils.py
+megatron/legacy/indexer.py
+megatron/training/__init__.py
+megatron/training/activations.py
+megatron/training/arguments.py
+megatron/training/async_utils.py
+megatron/training/checkpointing.py
+megatron/training/checkpointing_org.py
+megatron/training/dist_signal_handler.py
+megatron/training/ft_integration.py
+megatron/training/global_vars.py
+megatron/training/initialize.py
+megatron/training/log_handler.py
+megatron/training/one_logger_utils.py
+megatron/training/theoretical_memory_usage.py
+megatron/training/training.py
+megatron/training/utils.py
+megatron/training/wandb_utils.py
+megatron/training/yaml_arguments.py
+megatron_core.egg-info/PKG-INFO
+megatron_core.egg-info/SOURCES.txt
+megatron_core.egg-info/dependency_links.txt
+megatron_core.egg-info/requires.txt
+megatron_core.egg-info/top_level.txt
+requirements/pytorch_24.01/requirements.txt
+requirements/pytorch_24.07/requirements.txt
+requirements/pytorch_24.10/requirements.txt
\ No newline at end of file
diff --git a/aiak_megatron/megatron_core.egg-info/dependency_links.txt b/aiak_megatron/megatron_core.egg-info/dependency_links.txt
new file mode 100644
index 00000000..8b137891
--- /dev/null
+++ b/aiak_megatron/megatron_core.egg-info/dependency_links.txt
@@ -0,0 +1 @@
+
diff --git a/aiak_megatron/megatron_core.egg-info/requires.txt b/aiak_megatron/megatron_core.egg-info/requires.txt
new file mode 100644
index 00000000..9bf8fc9f
--- /dev/null
+++ b/aiak_megatron/megatron_core.egg-info/requires.txt
@@ -0,0 +1,21 @@
+einops
+flask-restful
+nltk
+pytest
+pytest_asyncio
+pytest-cov
+pytest_mock
+pytest-random-order
+sentencepiece
+tiktoken
+wrapt
+zarr
+wandb
+tensorstore!=0.1.46,!=0.1.72
+torch
+torch
+packaging
+
+[:sys_platform != "darwin"]
+nvidia-modelopt[torch]>=0.19.0
+nvidia-resiliency-ext
diff --git a/aiak_megatron/megatron_core.egg-info/top_level.txt b/aiak_megatron/megatron_core.egg-info/top_level.txt
new file mode 100644
index 00000000..3bfbb11f
--- /dev/null
+++ b/aiak_megatron/megatron_core.egg-info/top_level.txt
@@ -0,0 +1 @@
+megatron
diff --git a/aiak_megatron/pretrain_vlm.py b/aiak_megatron/pretrain_vlm.py
index f5cdbf33..a169d791 100644
--- a/aiak_megatron/pretrain_vlm.py
+++ b/aiak_megatron/pretrain_vlm.py
@@ -129,7 +129,10 @@ def model_provider(
# TODO: Make these configurable via input .yaml config.
vision_transformer_config = deepcopy(language_transformer_config)
- vision_transformer_config.num_layers = args.encoder_num_layers
+ # Only set num_layers from args for ViT models (like Qwen2-VL), not for FastViT
+ # FastViT determines its own layer count from architecture config
+ if not getattr(args, 'use_fastvit', False):
+ vision_transformer_config.num_layers = args.encoder_num_layers
vision_transformer_config.first_pipeline_num_layers = None
vision_transformer_config.last_pipeline_num_layers = None
vision_transformer_config.vision_model_type = vision_model_type
diff --git a/aiak_megatron/setup.py b/aiak_megatron/setup.py
index 4f8d121c..c00cd7bc 100644
--- a/aiak_megatron/setup.py
+++ b/aiak_megatron/setup.py
@@ -105,7 +105,7 @@ def req_file(filename, folder="requirements"):
'Natural Language :: English',
'Operating System :: OS Independent',
],
- packages=setuptools.find_namespace_packages(include=["megatron.core", "megatron.core.*"]),
+ packages=setuptools.find_namespace_packages(include=["megatron.core", "megatron.core.*","megatron.training","megatron.legacy"]),
ext_modules=[
Extension(
"megatron.core.datasets.helpers_cpp",
diff --git a/aiak_training_llm/data/chat_templete.py b/aiak_training_llm/data/chat_templete.py
index 82262cc4..136aee3f 100644
--- a/aiak_training_llm/data/chat_templete.py
+++ b/aiak_training_llm/data/chat_templete.py
@@ -98,6 +98,7 @@ class ChatTemplate:
efficient_eos: bool = False
replace_eos: bool = False
mm_plugin: Optional[MMPlugin] = None
+ use_tokenizer_template: bool = False # If True, use tokenizer's built-in chat template
def __post_init__(self):
if self.format_user is None:
@@ -163,6 +164,11 @@ def _encode(
Turn 0: prefix + system + query resp
Turn t: sep + query resp
"""
+ # Use tokenizer's built-in chat template if enabled
+ if self.use_tokenizer_template:
+ return self._encode_with_tokenizer_template(tokenizer, messages, system)
+
+ # Original custom template encoding
system = system or self.default_system
encoded_messages = []
for i, message in enumerate(messages):
@@ -187,6 +193,58 @@ def _encode(
return encoded_messages
+ def _encode_with_tokenizer_template(
+ self,
+ tokenizer: "AutoTokenizerFromHF",
+ messages: Sequence[Dict[str, str]],
+ system: Optional[str],
+ ) -> List[List[int]]:
+ """
+ Use tokenizer's built-in apply_chat_template for encoding.
+ This is useful for models like MobileLLM-R1 that have their own chat template.
+ """
+ # Prepare messages with system prompt if provided
+ formatted_messages = []
+ if system:
+ formatted_messages.append({"role": "system", "content": system})
+
+ # Add all user/assistant messages
+ for msg in messages:
+ formatted_messages.append(msg)
+
+ # Apply tokenizer's chat template
+ # Split into pairs: (user+system, assistant), (user, assistant), ...
+ encoded_messages = []
+
+ # Process each turn separately to get proper tokenization per turn
+ for i in range(0, len(messages), 2):
+ # User message (with system if first turn)
+ user_msgs = []
+ if i == 0 and system:
+ user_msgs.append({"role": "system", "content": system})
+ user_msgs.append(messages[i])
+
+ # Tokenize user turn
+ user_text = tokenizer.apply_chat_template(
+ user_msgs,
+ tokenize=False,
+ add_generation_prompt=True # Add assistant prompt
+ )
+ user_ids = tokenizer.encode(user_text, add_special_tokens=False)
+ encoded_messages.append(user_ids)
+
+ # Assistant message
+ if i + 1 < len(messages):
+ # Just the content, no special formatting
+ assistant_text = messages[i + 1]["content"]
+ # Add EOS if not efficient_eos
+ if not self.efficient_eos:
+ assistant_text += tokenizer.eos_token
+ assistant_ids = tokenizer.encode(assistant_text, add_special_tokens=False)
+ encoded_messages.append(assistant_ids)
+
+ return encoded_messages
+
def _convert_elements_to_ids(
self,
tokenizer: "AutoTokenizerFromHF",
@@ -279,6 +337,7 @@ def _register_chat_template(
efficient_eos: bool = False,
replace_eos: bool = False,
mm_plugin: Optional[MMPlugin] = None,
+ use_tokenizer_template: bool = False, # New parameter for auto template
) -> None:
"""
Registers a chat template.
@@ -320,6 +379,7 @@ def _register_chat_template(
efficient_eos=efficient_eos,
replace_eos=replace_eos,
mm_plugin=mm_plugin,
+ use_tokenizer_template=use_tokenizer_template, # Pass the new parameter
)
@@ -459,3 +519,13 @@ def get_support_templates() -> List[str]:
format_user=StringFormatter(slots=["<๏ฝUser๏ฝ>{{content}}<๏ฝAssistant๏ฝ>"]),
format_prefix=EmptyFormatter(slots=[{"bos_token"}]),
)
+
+
+# Auto template - uses tokenizer's built-in chat template
+# Perfect for MobileLLM-R1 and other models with embedded templates
+_register_chat_template(
+ name="auto",
+ cls=ChatTemplate,
+ use_tokenizer_template=True,
+ efficient_eos=True,
+)
diff --git a/aiak_training_llm/data/multimodal/dataloader_provider.py b/aiak_training_llm/data/multimodal/dataloader_provider.py
index ba27ac4a..5e781566 100644
--- a/aiak_training_llm/data/multimodal/dataloader_provider.py
+++ b/aiak_training_llm/data/multimodal/dataloader_provider.py
@@ -21,16 +21,21 @@ def get_train_dataset(task_encoder):
worker_debug_path=None,
worker_log_level=0
)
+ #use megatron energon to get the training dataset
+ # a high performance data loading library for large-scale distributed training
+ #each gpu gets its shard of the data based on data parallel rank
+ # Supports packing: Multiple short sequences can be packed into one to maximize GPU utilization
+ # streaming: can load data on-the-fly without pre-downloading entire dataset into memory
train_ds = energon.get_train_dataset(
- args.data_path[0],
- batch_size=args.micro_batch_size,
- task_encoder=task_encoder,
- worker_config=worker_config,
+ args.data_path[0], # Path to dataset (LLaVA-558K-Webdataset)
+ batch_size=args.micro_batch_size, #Batch size per GPU
+ task_encoder=task_encoder, # The Qwen2VLTaskEncoder we just created
+ worker_config=worker_config, #Distributed Config
max_samples_per_sequence=None,
shuffle_buffer_size=None,
- packing_buffer_size=args.packing_batch_size,
- handler=print_error_handler,
- image_decode="pil",
+ packing_buffer_size=args.packing_batch_size, #Buffer size for packing sequences
+ handler=print_error_handler, #Error handler function
+ image_decode="pil", #Decode images using PIL
)
return train_ds
@@ -65,11 +70,12 @@ class EnergonDataloader:
def __init__(self, dataloader, collator=None):
self._dataloader = dataloader
self._collator = collator
- self._iter = iter(cyclic_iter(dataloader))
+ self._iter = iter(cyclic_iter(dataloader)) # Infinite iterator!
def __next__(self):
- features = self._iter.__next__()
+ features = self._iter.__next__() #get next batch
if self._collator is not None:
+ # Apply padding using the collator
padded = self._collator.tokenizer.pad(
{"input_ids": features['tokens']},
padding=self._collator.padding,
@@ -77,7 +83,10 @@ def __next__(self):
pad_to_multiple_of=self._collator.pad_to_multiple_of,
)
paded_length = padded['input_ids'].shape[1] - features['tokens'].shape[1]
+
+ # Update features with padded versions
features['tokens'] = padded["input_ids"]
+ #pad labels and attention mask too
features['labels'] = F.pad(
features['labels'],
(0, paded_length),
@@ -85,6 +94,11 @@ def __next__(self):
self._collator.label_pad_token_id
)
features['attn_mask'] = F.pad(features['attn_mask'], (0, paded_length), "constant", True)
+ if os.getenv("VERBOSE_DATALOADER_DEBUG", "0").lower() in {"1", "true", "yes", "on"}:
+ print("Dataloader batch - tokens shape:", features['tokens'].shape)
+ print("Dataloader batch - labels shape:", features['labels'].shape)
+ print("Dataloader batch - attn_mask shape:", features['attn_mask'].shape)
+ print("Dataloader batch - image_grid_thw shape:", features['image_grid_thw'].shape)
return features
def __iter__(self):
@@ -94,7 +108,7 @@ def save_state(self):
""" Save the current state of this dataloader """
return self._dataloader.save_state_rank()
-
+# This means training never runs out of data - it just loops back to the beginning.
def cyclic_iter(iter):
""" Infinite iteration over an iterator """
while True:
diff --git a/aiak_training_llm/data/multimodal/qwen2vl_task_encoder.py b/aiak_training_llm/data/multimodal/qwen2vl_task_encoder.py
index 9b30a33b..39684eaa 100644
--- a/aiak_training_llm/data/multimodal/qwen2vl_task_encoder.py
+++ b/aiak_training_llm/data/multimodal/qwen2vl_task_encoder.py
@@ -1,4 +1,6 @@
""" Qwen2VLTaskEncoder class."""
+# The Qwen2VLTaskEncoder is responsible for converting raw data samples (images + text conversations) into model-ready tensors. It's called by Megatron Energon during data loading.
+
import math
import re
from dataclasses import dataclass, field
@@ -15,14 +17,30 @@
from torchvision.transforms import InterpolationMode
from transformers import AutoProcessor
from typing_extensions import override
+#Fast ViT imports
+# Custom Image processor for FastViT
+from aiak_training_llm.models.fastvit.fastvit_preprocessor import FastViTImageProcessor
+# Utilities for FastViT multimodal image preprocessing
+from aiak_training_llm.models.fastvit.mm_utils import (
+ expand2square, # pads image to square with background color
+ process_anyres_image, # handles variable resolution with patches
+ process_images, #Main dispatcher based on aspect ratio mode
+)
from aiak_training_llm.data.multimodal import MultiMixQASample
from aiak_training_llm.data.multimodal.length_sort_dataset import LengthPoolSortDataset
-from aiak_training_llm.utils import constants, get_chat_template
+from aiak_training_llm.utils import constants, get_chat_template, get_tokenizer
+from aiak_training_llm.tokenizer.special_tokens import ensure_multimodal_special_tokens
from .task_encoder import (ImageTaskBatchPacked, ImageTaskSample,
ImageTaskSamplePacked, TaskEncoder)
+# Print helper: only emit from rank-0 process so multi-GPU runs don't duplicate output.
+import os as _os
+def _rank0_print(*args, **kwargs):
+ if _os.environ.get("RANK", "0") == "0":
+ print(*args, **kwargs)
+
from megatron.energon.flavors.base_dataset import (
BaseCoreDatasetFactory,
PinMemoryMixin,
@@ -112,22 +130,96 @@ class Qwen2VLTaskEncoder(TaskEncoder):
def __init__(self, args):
super().__init__()
+ self.args = args
if args.training_phase in ['sft']:
- self.chat_template = get_chat_template()
- self.processor = AutoProcessor.from_pretrained(self.args.hf_tokenizer_path, trust_remote_code=True)
+ self.chat_template = get_chat_template() # Load conversation template
+ #template for formatting conversations (user/assistant turns)
+
+ # OLD APPROACH (commented out - didn't work for MobileLLM):
+ # Load HuggingFace processor (tokenizer + image/ video processor from Qwen2-VL)
+ # self.processor = AutoProcessor.from_pretrained(self.args.hf_tokenizer_path, trust_remote_code=True, local_files_only=False)
+ # print(f"Loaded processor from {self.args.hf_tokenizer_path}")
+ # print(f"Processor config: {self.processor}")
+
+ # NEW APPROACH: Handle FastVLM (MobileLLM + FastViT) vs Qwen2-VL separately
+ # FastViT image processor (following FastVLM repo)
+ # Use this for vision encoding instead of Qwen2-VL's processor
+ self.use_fastvit = getattr(args, 'use_fastvit', False)
+
+ if self.use_fastvit:
+ # For FastVLM: Load only the language model tokenizer (e.g., MobileLLM)
+ # No need for a full processor since FastViT handles vision separately
+ from transformers import AutoTokenizer
+ shared_tokenizer = get_tokenizer()
+ if shared_tokenizer is not None and hasattr(shared_tokenizer, "hf_tokenizer"):
+ self.tokenizer_hf = shared_tokenizer.hf_tokenizer()
+ _rank0_print("Using shared training tokenizer in FastVLM mode")
+ else:
+ self.tokenizer_hf = AutoTokenizer.from_pretrained(
+ self.args.hf_tokenizer_path,
+ trust_remote_code=True,
+ local_files_only=False,
+ )
+
+ added_mm_tokens = ensure_multimodal_special_tokens(self.tokenizer_hf)
+ _rank0_print(f"Ensured multimodal special tokens for FastVLM tokenizer; added={added_mm_tokens}")
+
+ # Set padding token if not present (required for batch processing)
+ if self.tokenizer_hf.pad_token is None:
+ self.tokenizer_hf.pad_token = self.tokenizer_hf.eos_token
+ _rank0_print(f"Set pad_token to eos_token: {self.tokenizer_hf.eos_token}")
+
+ # Create a wrapper that properly delegates all methods to the tokenizer
+ class TokenizerWrapper:
+ def __init__(self, tokenizer):
+ self.tokenizer = tokenizer
+
+ def __getattr__(self, name):
+ # Delegate all other attributes/methods to the underlying tokenizer
+ return getattr(self.tokenizer, name)
+
+ self.processor = TokenizerWrapper(self.tokenizer_hf)
+ _rank0_print(f"Loaded tokenizer (FastVLM mode) from {self.args.hf_tokenizer_path}")
+
+ # Initialize FastViT image processor
+ fastvit_image_size = getattr(args, 'fastvit_image_size', 1024)
+ self.fastvit_processor = FastViTImageProcessor(image_size=fastvit_image_size)
+ fastvit_patch_size = getattr(args, 'fastvit_patch_size', 64)
+ self.fastvit_tokens_per_side = fastvit_image_size // fastvit_patch_size
+ self.fastvit_num_image_tokens = self.fastvit_tokens_per_side ** 2
+ self.image_token_block = IMAGE_TOKEN * self.fastvit_num_image_tokens
+ _rank0_print(
+ f"Initialized FastViT processor with image_size={fastvit_image_size}, "
+ f"tokens_per_image={self.fastvit_num_image_tokens}"
+ )
+ else:
+ # For Qwen2-VL: Load full processor (tokenizer + image/video processor)
+ self.processor = AutoProcessor.from_pretrained(
+ self.args.hf_tokenizer_path,
+ trust_remote_code=True,
+ local_files_only=False
+ )
+ _rank0_print(f"Loaded Qwen2-VL processor from {self.args.hf_tokenizer_path}")
+ self.image_token_block = IMAGE_TOKEN_WITH_TAGS
+ _rank0_print(f"Processor config: {self.processor}")
+
+ #Resolution parameters for resizing images/videos
if args.image_resolution:
setattr(self.processor, 'image_resolution', args.image_resolution)
- # video
+ # resolution parameters for resizing images/videos
+ _rank0_print("image_resolution:", getattr(self.processor, 'image_resolution', None))
+ # Video processing parameters
self.frame_min_pixels = args.frame_min_pixels
self.frame_max_pixels = args.frame_max_pixels
self.video_max_pixels = args.video_max_pixels
self.fps = args.fps
self.fps_min_frames = args.fps_min_frames
self.fps_max_frames = args.fps_max_frames
- # image
+ # Image processing parameters
self.min_pixels = args.min_pixels
self.max_pixels = args.max_pixels
+ self._logged_multimodal_token_debug_once = False
def _reisize_video(self, vision: VideoData, image_factor=28, frame_factor=2):
""" Resize video: frame number, height, width """
@@ -160,19 +252,121 @@ def _reisize_video(self, vision: VideoData, image_factor=28, frame_factor=2):
return video
def _resize_image(self, image, size_factor=28):
+ """
+ Resize image based on vision encoder type.
+ - For FastViT: Use FastVLM's preprocessing (pad/anyres)
+ - For Rice/SigLIP: Dynamic resize with smart_resize
+ """
+ if self.use_fastvit:
+ # FastViT: Use FastVLM's preprocessing approach
+ # For single image, we'll handle aspect ratio in _process
+ # Just return the PIL image as-is for now
+ return image
+
+ # Original Rice/SigLIP preprocessing
+ #calculate optimal size using smart_resize
+ # constraints:
+ # 1- width and height must be multiple of size_factor (28)
+ # 2- Total pixels (height ร width) are โฅ min_pixels
+ # 3- Total pixels (height ร width) are โค max_pixels
+ # 4- Aspect ratio is preserved as much as possible
+
resized_height, resized_width = smart_resize(
image.height,
image.width,
- factor=size_factor,
- min_pixels=self.min_pixels,
- max_pixels=self.max_pixels,
+ factor=size_factor,
+ # why factor of 28?
+ # Input Image
+ # โ
+ # Split into 14ร14 patches (Patch Embedding)
+ # โ
+ # 2ร2 patch merging (Reduce spatial dimensions by 2)
+ # โ
+ # Effective patch size = 14 ร 2 = 28ร28 pixels
+ # Example:
+
+ # Image size: 1120 ร 784 pixels
+ # Number of patches: (1120/28) ร (784/28) = 40 ร 28 = 1,120 patches
+ # Each patch โ 1 vision token
+ min_pixels=self.min_pixels, # e.g., 256*28*28
+ max_pixels=self.max_pixels, # e.g., 1280*28*28
)
image = image.resize((resized_width, resized_height))
- return image
+ return image # return resized PIL image
def _process(self, image, text):
"""" Process the data to get the model's input """
+ if self.use_fastvit and image is not None:
+ # FastViT preprocessing using FastVLM's approach
+ # Tokenize text only
+ text_inputs = self.processor.tokenizer(
+ text=text,
+ padding=True,
+ return_tensors="pt",
+ )
+ input_ids = text_inputs['input_ids'][0]
+ attn_mask = text_inputs['attention_mask'][0].logical_not()
+ fastvit_tokenizer = self.processor.tokenizer
+ vision_start_id, img_pad_id, vision_end_id = fastvit_tokenizer.convert_tokens_to_ids([
+ VISION_TAGS[0],
+ IMAGE_TOKEN,
+ VISION_TAGS[1]
+ ])
+ if vision_start_id is None or img_pad_id is None or vision_end_id is None:
+ vocab = fastvit_tokenizer.get_vocab()
+ if img_pad_id is None:
+ img_pad_id = vocab.get(IMAGE_TOKEN)
+ if vision_start_id is None:
+ vision_start_id = vocab.get(VISION_TAGS[0])
+ if vision_end_id is None:
+ vision_end_id = vocab.get(VISION_TAGS[1])
+ if not self._logged_multimodal_token_debug_once:
+ img_count = int((input_ids == img_pad_id).sum().item()) if img_pad_id is not None else 0
+ vstart_count = int((input_ids == vision_start_id).sum().item()) if vision_start_id is not None else 0
+ _rank0_print(
+ f"[DEBUG PREPROCESS TOKENS] image_token='{IMAGE_TOKEN}' id={img_pad_id}, "
+ f"vision_start='{VISION_TAGS[0]}' id={vision_start_id}, vision_end='{VISION_TAGS[1]}' id={vision_end_id}, "
+ f"counts_in_input_ids: image={img_count}, vision_start={vstart_count}"
+ )
+ self._logged_multimodal_token_debug_once = True
+
+ # Process image with FastVLM's preprocessing utilities
+ # Default to 'pad' aspect ratio (expand to square with padding)
+ image_aspect_ratio = getattr(self.args, 'image_aspect_ratio', 'pad')
+ if image_aspect_ratio == 'pad':
+ # Expand to square with mean color padding
+ mean_color = tuple(int(x * 255) for x in self.fastvit_processor.image_mean)
+ image = expand2square(image, mean_color)
+ pixel_values = self.fastvit_processor(image)
+
+
+ elif image_aspect_ratio == 'anyres':
+ # Process with variable resolution (patches)
+ grid_pinpoints = getattr(self.args, 'image_grid_pinpoints', '[(384, 384), (768, 384), (384, 768), (768, 768)]')
+ pixel_values = process_anyres_image(
+ image,
+ self.fastvit_processor.processor, # CLIPImageProcessor
+ grid_pinpoints
+ )
+ else:
+ # Direct resize to target size
+ pixel_values = self.fastvit_processor(image)
+
+ pixel = [pixel_values]
+
+ # FastViT: one square tile with a spatial token grid.
+ image_grid_thw = torch.tensor([[1, self.fastvit_tokens_per_side, self.fastvit_tokens_per_side]])
+
+ # Create target tensor (same as Qwen2-VL path)
+ target = input_ids.clone()
+ for token_id in (vision_start_id, img_pad_id, vision_end_id):
+ if token_id is not None:
+ target[target == token_id] = IGNORE_INDEX
+
+ return input_ids, target, pixel, image_grid_thw, attn_mask
+
+ # Original Qwen2-VL processing
inputs = self.processor(
text=text,
images=image,
@@ -193,6 +387,15 @@ def _process(self, image, text):
IMAGE_TOKEN,
VISION_TAGS[1]
])
+ if not self._logged_multimodal_token_debug_once:
+ img_count = int((input_ids == img_pad_id).sum().item()) if img_pad_id is not None else 0
+ vstart_count = int((input_ids == vision_start_id).sum().item()) if vision_start_id is not None else 0
+ _rank0_print(
+ f"[DEBUG PREPROCESS TOKENS] image_token='{IMAGE_TOKEN}' id={img_pad_id}, "
+ f"vision_start='{VISION_TAGS[0]}' id={vision_start_id}, vision_end='{VISION_TAGS[1]}' id={vision_end_id}, "
+ f"counts_in_input_ids: image={img_count}, vision_start={vstart_count}"
+ )
+ self._logged_multimodal_token_debug_once = True
target[target == vision_start_id] = IGNORE_INDEX
target[target == img_pad_id] = IGNORE_INDEX
target[target == vision_end_id] = IGNORE_INDEX
@@ -212,7 +415,7 @@ def process_sft_vqa(self, context, answer, image):
}],
tokenize=False
).replace(
- "", IMAGE_TOKEN_WITH_TAGS
+ "", self.image_token_block
)
if text[-1] == '\n':
text = text[:-1]
@@ -226,6 +429,12 @@ def process_sft_vqa(self, context, answer, image):
def process_sft_qa(self, messages: list, system: str, raw_video: list, raw_image: list):
""" process the data for sft qa """
+ # messages: List of conversation turns (user/assistant)
+ # system: System prompt
+ # raw_video: List of PIL/VideoData objects (or None)
+ # raw_image: List of PIL.Image objects (or None)
+
+ #Initialize output containers
video_grid_thw = None
pixel_values_videos = []
image_grid_thw = None
@@ -233,22 +442,43 @@ def process_sft_qa(self, messages: list, system: str, raw_video: list, raw_image
video = []
image = []
-
+ # resize image
if raw_image is not None:
+ # loop through each image and resize
for i in raw_image:
image.append(self._resize_image(i))
-
+ # resize video
if raw_video is not None:
for v in raw_video:
video.append(self._reisize_video(v))
+
+ # input messages:
+ # [
+ # {"from": "human", "value": "\nWhat's in this image?"},
+ # {"from": "gpt", "value": "A cat sitting on a mat."}
+ # ]
messages, mm_inputs = self.chat_template.mm_plugin.process_messages(
messages,
image if image is not None else [],
video if raw_video is not None else [],
self.processor
)
+ # output messages:
+ # [
+ # {"from": "human", "value": "<|vision_start|><|image_pad|><|image_pad|>...<|vision_end|>\nWhat's in this image?"},
+ # {"from": "gpt", "value": "A cat sitting on a mat."}
+ # ]
+ # Output mm_inputs:
+ # {
+ # "pixel_values": Tensor([1, 1176, 1176]), # Vision encoder input
+ # "image_grid_thw": Tensor([[1, 28, 42]]), # 1 temporal, 28 height patches, 42 width patches
+ # }
+
+
# assert raw_image is not None, f'No image found in {messages}' ็กฎๅฎๆ็บฏๆๆฌๅฏน่ฏ
+
+ #extracting the multi-modal inputs
if raw_video is not None:
video_grid_thw = mm_inputs["video_grid_thw"]
pixel_values_videos = [mm_inputs["pixel_values_videos"]]
@@ -257,21 +487,71 @@ def process_sft_qa(self, messages: list, system: str, raw_video: list, raw_image
pixel_values_images = [mm_inputs["pixel_values"]]
encode_pairs = self.chat_template.encode_multiturn(
+ # 1. Applies chat template to format conversation:
+ # example:
+ # <|im_start|>system
+ # You are a helpful assistant.<|im_end|>
+ # <|im_start|>user
+ # <|vision_start|><|image_pad|><|image_pad|>...<|vision_end|>
+ # What's in this image?<|im_end|>
+ # <|im_start|>assistant
+ # A cat sitting on a mat.<|im_end|>
+
+ # 2. Tokenizes the entire formatted conversation
+
+ # 3. Splits into (source, target) pairs for each turn:
+ # example:
+ # encode_pairs = [
+ # (
+ # [151644, 8948, 198, ...], # source_ids: system + user prompt (don't train)
+ # [8122, 4758, 6134, ...] # target_ids: assistant response (train on this)
+ # )
+ # ]
+
+
tokenizer=self.tokenizer,
messages=messages,
system=system,
)
+
input_ids, target = [], []
for turn_idx, (source_ids, target_ids) in enumerate(encode_pairs):
- input_ids += source_ids + target_ids
+ input_ids += source_ids + target_ids # Append both source and target to input_ids
+ # input_ids = [151644, 8948, ..., 151645, 8122, 4758, ...]
+ # ^^^^^^^^^^^^^^^^^^^^^ ^^^^^^^^^^^^^^^^^^^
+ # system + user assistant response
+
+ # Mask source (user prompt), keep target (assistant response)
target += [IGNORE_INDEX] * len(source_ids) + target_ids
- input_ids = torch.tensor(input_ids)
+ # target = [-100, -100, ..., -100, 8122, 4758, ...]
+ # ^^^^^^^^^^^^^^^^^^^^^^^^^ ^^^^^^^^^^^^^^^^^^^
+ # Masked (don't train on) Train on this!
+ # Why mask source_ids?
+ # We don't want the model to learn to predict the user's input
+ # only the assistant's response!
+
+ # Convert to Tensors and Create Attention Mask
+ input_ids = torch.tensor(input_ids) # Shape: [seq_len]
+ # input_ids: Tensor([151644, 8948, 198, ..., 8122, 4758, ...])
target = torch.tensor(target)
+ # target: Tensor([-100, -100, ..., 8122, 4758, ...])
+
+ # Create attention mask (all False = attend to all tokens)
attn_mask = torch.zeros_like(input_ids).bool()
+ # attn_mask: Tensor([False, False, False, ..., False, False]) # Shape: [seq_len]
+
return input_ids, target, attn_mask, pixel_values_images, image_grid_thw, \
pixel_values_videos, video_grid_thw
+ # input_ids: Tokenized conversation - Shape: [seq_len]
+ # target: Labels with masking - Shape: [seq_len]
+ # attn_mask: Attention mask - Shape: [seq_len]
+ # pixel_values_images: Image tensors - List of Tensors
+ # image_grid_thw: Image grid info - Tensor [[1, 28, 42]]
+ # pixel_values_videos: Video tensors - Empty list
+ # video_grid_thw: Video grid info - None
+
def encode_captioning(self, sample: CaptioningSample) -> ImageTaskSample:
"""Encode CaptioningSample."""
@@ -282,14 +562,15 @@ def encode_captioning(self, sample: CaptioningSample) -> ImageTaskSample:
# assert self.args.training_phase == constants.TrainingPhase.PRETRAIN, "Only support PRETRAIN phase"
- text = IMAGE_TOKEN_WITH_TAGS + sample.caption + self.tokenizer.tokenizer.eos_token
+ # text = self.image_token_block + sample.caption + self.tokenizer.tokenizer.eos_token
+ text = self.image_token_block + sample.caption + self.tokenizer.eos_token
input_ids, target, imgs, image_grid_thw, attn_mask = self._process(sample.image, text)
- num_tiles = [len(image_grid_thw)]
+ num_tiles = [len(image_grid_thw)] if image_grid_thw is not None else [1]
if self.args.enable_discard_sample:
assert len(input_ids) <= self.args.seq_length, f"{sample.__key__} input length {len(input_ids)}"
- else:
+ elif image_grid_thw is not None:
assert image_grid_thw.prod() / 4 <= self.args.seq_length, f"{sample.__key__} thw {image_grid_thw}"
return Qwen2VLImageTaskSample(
@@ -320,12 +601,11 @@ def encode_vqa4packing(self, sample: VQASample) -> ImageTaskSample:
'content': sample.answers
}],
tokenize=False
- ).replace("", IMAGE_TOKEN_WITH_TAGS)
+ ).replace("", self.image_token_block)
if text[-1] == '\n':
text = text[:-1]
pass
-
input_ids, _, imgs, image_grid_thw, attn_mask = self._process(sample.image, text)
target = torch.ones_like(input_ids) * IGNORE_INDEX
answers = self.tokenizer.tokenize(sample.answers)
@@ -333,10 +613,10 @@ def encode_vqa4packing(self, sample: VQASample) -> ImageTaskSample:
target[-1] = input_ids[-1]
# print(target[-1])
- num_tiles = [len(image_grid_thw)]
+ num_tiles = [len(image_grid_thw)] if image_grid_thw is not None else [1]
if self.args.enable_discard_sample:
assert len(input_ids) <= self.args.seq_length, f"{sample.__key__} input length {len(input_ids)}"
- else:
+ elif image_grid_thw is not None:
assert image_grid_thw.prod() / 4 <= self.args.seq_length, f"{sample.__key__} grid_thw: {image_grid_thw}"
return Qwen2VLImageTaskSample(
@@ -357,6 +637,7 @@ def encode_vqa4packing(self, sample: VQASample) -> ImageTaskSample:
def encode_multi_vid_qa(self, sample: VQASample) -> ImageTaskSample:
"""Encode sample in Qwen2VL style."""
if self.args.training_phase == constants.TrainingPhase.SFT:
+ # call main processing function process_sft_qa
input_ids, target, attn_mask, imgs, image_grid_thw, video, video_grid_thw = \
self.process_sft_qa(sample.messages, sample.system, sample.video, None)
else:
@@ -364,7 +645,7 @@ def encode_multi_vid_qa(self, sample: VQASample) -> ImageTaskSample:
if self.args.enable_discard_sample:
assert len(input_ids) <= self.args.seq_length, f"{sample.__key__} input length {len(input_ids)}"
- else:
+ elif video_grid_thw is not None:
assert video_grid_thw.prod(dim=-1).sum() / 4 <= self.args.seq_length, \
f"{sample.__key__} grid_thw: {video_grid_thw}"
@@ -384,18 +665,19 @@ def encode_multi_vid_qa(self, sample: VQASample) -> ImageTaskSample:
total_len=len(input_ids),
)
-
+ # Energon automatically calls the appropriate encode method based on sample type.
+ # For SFT with multi-modal data, it calls:
def encode_multi_mix_qa(self, sample: MultiMixQASample) -> ImageTaskSample:
"""Encode sample in Qwen2VL style."""
if self.args.training_phase == constants.TrainingPhase.SFT:
- num_tiles = []
-
+ num_tiles = [] #store number of tiles for each image/ video after processing
+ # call main processing function process_sft_qa
input_ids, target, attn_mask, imgs, image_grid_thw, pixel_values_videos, video_grid_thw = \
self.process_sft_qa(sample.messages, sample.system, sample.video, sample.image)
if sample.video is not None:
- num_tiles = [len(video_grid_thw)]
+ num_tiles = [len(video_grid_thw)] if video_grid_thw is not None else [1]
elif sample.image is not None:
- num_tiles = [len(image_grid_thw)]
+ num_tiles = [len(image_grid_thw)] if image_grid_thw is not None else [1]
else:
raise NotImplementedError(f"Unknown training phase {self.args.training_phase}")
@@ -405,10 +687,10 @@ def encode_multi_mix_qa(self, sample: MultiMixQASample) -> ImageTaskSample:
if self.args.enable_discard_sample:
assert len(input_ids) <= self.args.seq_length, f"{sample.__key__} input length {len(input_ids)}"
- elif sample.video is not None:
+ elif sample.video is not None and video_grid_thw is not None:
assert video_grid_thw.prod(dim=-1).sum() / 4 <= self.args.seq_length, \
f"{sample.__key__} grid_thw: {video_grid_thw}"
- elif sample.image is not None:
+ elif sample.image is not None and image_grid_thw is not None:
assert image_grid_thw.prod(dim=-1).sum() / 4 <= self.args.seq_length, \
f"{sample.__key__} grid_thw: {image_grid_thw}"
@@ -435,11 +717,12 @@ def encode_vaq(self, sample: VQASample) -> ImageTaskSample:
if self.args.add_question_in_pretrain:
text = (sample.context + sample.answers).replace(
"",
- IMAGE_TOKEN_WITH_TAGS
+ self.image_token_block
)
else:
- text = IMAGE_TOKEN_WITH_TAGS + sample.answers
- text = text + self.tokenizer.tokenizer.eos_token
+ text = self.image_token_block + sample.answers
+ # text = text + self.tokenizer.tokenizer.eos_token
+ text = text + self.tokenizer.eos_token
input_ids, target, imgs, image_grid_thw, attn_mask = self._process(sample.image, text)
elif self.args.training_phase == constants.TrainingPhase.SFT:
@@ -481,7 +764,7 @@ def encode_vaq(self, sample: VQASample) -> ImageTaskSample:
# Fallback to a hard cut of the original preliminary string if no sentence ender is found.
sample.answers = preliminary_cut
- print(
+ _rank0_print(
f"Answer truncated to a full sentence. "
f"Original length: {original_length}, New length: {len(sample.answers)}"
)
@@ -495,7 +778,7 @@ def encode_vaq(self, sample: VQASample) -> ImageTaskSample:
'content': sample.answers
}],
tokenize=False
- ).replace("", IMAGE_TOKEN_WITH_TAGS)
+ ).replace("", self.image_token_block)
if text[-1] == '\n':
text = text[:-1]
input_ids, _, imgs, image_grid_thw, attn_mask = self._process(sample.image, text)
@@ -507,7 +790,7 @@ def encode_vaq(self, sample: VQASample) -> ImageTaskSample:
else:
raise NotImplementedError(f"Unknown training phase {self.args.training_phase}")
- num_tiles = [len(image_grid_thw)]
+ num_tiles = [len(image_grid_thw)] if image_grid_thw is not None else [1]
if self.args.enable_discard_sample:
assert len(input_ids) <= self.args.seq_length, f"{sample.__key__} input length {len(input_ids)}"
@@ -557,12 +840,18 @@ def pack_selected_samples(self, samples: List[Qwen2VLImageTaskSample]) -> List[Q
)
@override
+ # After encoding , multiple encoded samples are batched together
def batch(self, samples: List[Union[Qwen2VLImageTaskSample, Qwen2VLImageTaskSamplePacked]]) \
-> Qwen2VLImageTaskBatchPacked:
""" Batch samples together """
image_grid_thw, video_grid_thw = self.process_samples_grid(samples)
+ # Create batch
return Qwen2VLImageTaskBatchPacked(
- super().batch(samples),
+ super().batch(samples), # # Stacks tokens, labels, etc. with padding
+ # Pads tokens to same length โ [batch_size, max_seq_len]
+ # Pads labels to same length โ [batch_size, max_seq_len]
+ # Concatenates imgs โ [total_images_in_batch, C, H, W]
+ # Creates batch attention masks
image_grid_thw=image_grid_thw,
video_grid_thw=video_grid_thw
)
@@ -701,4 +990,4 @@ def build_train_datasets(
)
if worker_config.should_log(level=1):
dataset = LogSampleDataset(dataset, mode="train", worker_config=worker_config)
- return dataset
\ No newline at end of file
+ return dataset
diff --git a/aiak_training_llm/data/multimodal/task_encoder.py b/aiak_training_llm/data/multimodal/task_encoder.py
index ee3bd692..a7fe84a1 100644
--- a/aiak_training_llm/data/multimodal/task_encoder.py
+++ b/aiak_training_llm/data/multimodal/task_encoder.py
@@ -7,9 +7,14 @@
import os
import sys
import traceback
+
+def _rank0_print(*args, **kwargs):
+ """Print only from rank-0 process to avoid duplicated output in multi-GPU runs."""
+ if os.environ.get("RANK", "0") == "0":
+ print(*args, **kwargs)
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple, Union
-
+import logging
from PIL import Image
from torchvision.transforms import ToPILImage
import numpy as np
@@ -309,22 +314,239 @@ def select_samples_to_pack(self, samples: List[ImageTaskSample]) -> List[List[Im
return packed_samples
+ # @stateless
+ # def pack_selected_samples(self, samples: List[ImageTaskSample]) -> List[ImageTaskSamplePacked]:
+ # """
+ # Function to pack a list of ImageTaskSample into a single ImageTaskSamplePacked.
+
+ # NOTE: Energon dataloader calls this method internally if packing is used.
+ # Please see https://nvidia.github.io/Megatron-Energon/packing.html
+
+ # Args:
+ # samples: List of ImageTaskSample instances to pack into one sample.
+
+ # Returns:
+ # ImageTaskSamplePacked instance.
+ # """
+
+ # packing_seq_len = self.args.seq_length
+
+ # packed_tokens = []
+ # packed_labels = []
+ # packed_masks = []
+ # packed_imgs = []
+ # packed_videos = []
+
+ # current_length = 0
+ # max_length = 0
+ # cu_lengths = [0]
+
+ # # Process each sample and build lists that we will concatenate to create the packed sample.
+ # for _, sample in enumerate(samples):
+ # sample_len = sample.total_len
+
+ # if sample_len > max_length:
+ # max_length = sample_len
+
+
+ # # If total_len property is used, update it to reflect truncation
+ # sample_len = min(sample_len, packing_seq_len)
+ # try:
+ # sample.total_len = sample_len
+ # except Exception:
+ # # If dataclass is frozen or doesn't allow assign, ignore;
+ # # we will use sample_len variable for logic below.
+ # pass
+
+ # # update max_length to be the maximum sample length after truncation
+ # if sample_len > max_length:
+ # max_length = sample_len
+
+
+
+ # # If adding this sample exceeds the max length, stop.
+ # # This should not happen.
+ # # The select_samples_to_pack method should have already ensured that the samples fit.
+ # if current_length + sample_len > packing_seq_len:
+ # print(f"packing_seq_len:{packing_seq_len}----<<<<<----{current_length + sample_len}")
+ # raise ValueError(f"Packed sample exceeds the maximum sequence length of {packing_seq_len}: {samples}")
+
+ # # Add the sample's tokens and labels
+ # packed_tokens.append(sample.tokens)
+ # packed_labels.append(sample.labels)
+ # packed_masks.append(sample.attn_mask)
+
+ # # Add the images
+ # if sample.imgs is not None:
+ # packed_imgs += sample.imgs
+ # if sample.pixel_values_videos is not None:
+ # packed_videos += sample.pixel_values_videos
+ # current_length += sample_len
+ # cu_lengths.append(current_length)
+
+ # # Concatenate packed tokens and labels.
+ # packed_tokens = torch.cat(packed_tokens, dim=0)
+ # packed_labels = torch.cat(packed_labels, dim=0)
+ # packed_masks = torch.cat(packed_masks, dim=0)
+
+ # return ImageTaskSamplePacked(
+ # __key__=",".join([s.__key__ for s in samples]),
+ # __restore_key__=(), # Will be set by energon based on `samples`
+ # __subflavor__=None,
+ # __subflavors__=samples[0].__subflavors__,
+ # tokens=packed_tokens,
+ # labels=packed_labels,
+ # attn_mask=packed_masks,
+ # imgs=packed_imgs,
+ # pixel_values_videos=packed_videos,
+ # cu_lengths=torch.tensor(cu_lengths, dtype=torch.int32),
+ # max_length=max_length,
+ # num_tiles=[n for s in samples for n in s.num_tiles],
+ # )
+
+ # @stateless
+ # def pack_selected_samples(self, samples: List[ImageTaskSample]) -> ImageTaskSamplePacked:
+ # """
+ # Pack a list of ImageTaskSample into a single ImageTaskSamplePacked.
+ # Truncates samples as needed so the packed sample never exceeds packing_seq_len.
+ # """
+ # packing_seq_len = self.args.seq_length
+
+ # # Lists of per-token pieces (will be concatenated at the end).
+ # packed_tokens_parts = []
+ # packed_labels_parts = []
+ # packed_masks_parts = []
+ # packed_imgs = []
+ # packed_videos = []
+
+ # current_length = 0
+ # max_length = 0
+ # cu_lengths = [0]
+
+ # for sample in samples:
+ # # Determine sample token length (use available fields)
+ # sample_len = int(getattr(sample, "total_len", None) or
+ # (len(sample.tokens) if hasattr(sample, "tokens") and sample.tokens is not None else 0))
+
+ # # If the sample itself is longer than the global packing_seq_len, allow truncation
+ # if sample_len > packing_seq_len:
+ # logging.getLogger(__name__).warning(
+ # f"Sample {getattr(sample, '__key__', '')} length {sample_len} > packing_seq_len {packing_seq_len}. "
+ # "Truncating token fields to packing_seq_len."
+ # )
+ # sample_len = packing_seq_len
+
+ # remaining_capacity = packing_seq_len - current_length
+ # if remaining_capacity <= 0:
+ # # no capacity left (shouldn't normally happen if select_samples_to_pack worked),
+ # # stop adding further samples.
+ # logging.getLogger(__name__).warning(
+ # "No remaining packing capacity; stopping packing of further samples."
+ # )
+ # break
+
+ # # If this sample would overflow the remaining capacity, we will truncate it to fit.
+ # if sample_len > remaining_capacity:
+ # logging.getLogger(__name__).warning(
+ # f"Truncating sample {getattr(sample, '__key__', '')} from {sample_len} to fit remaining capacity {remaining_capacity}."
+ # )
+ # sample_len = remaining_capacity
+
+ # # Slice the per-token fields safely (they might be torch tensors or numpy arrays)
+ # # If a field is missing, create an appropriate dummy slice of length sample_len.
+ # def slice_field(field, length):
+ # if field is None:
+ # # return a zero-length tensor/array (we'll handle concatenation later)
+ # return None
+ # try:
+ # return field[:length]
+ # except Exception:
+ # # fallback: convert to torch tensor then slice
+ # try:
+ # t = torch.as_tensor(field)
+ # return t[:length]
+ # except Exception:
+ # return None
+
+ # tokens_slice = slice_field(getattr(sample, "tokens", None), sample_len)
+ # labels_slice = slice_field(getattr(sample, "labels", None), sample_len)
+ # mask_slice = slice_field(getattr(sample, "attn_mask", None), sample_len)
+
+ # # Append slices (if None, we will handle later by creating empty tensors)
+ # packed_tokens_parts.append(tokens_slice)
+ # packed_labels_parts.append(labels_slice)
+ # packed_masks_parts.append(mask_slice)
+
+ # # For images/videos: keep them as-is (we assume images are not token-aligned)
+ # if getattr(sample, "imgs", None) is not None:
+ # packed_imgs += list(sample.imgs)
+ # if getattr(sample, "pixel_values_videos", None) is not None:
+ # packed_videos += list(sample.pixel_values_videos)
+
+ # current_length += sample_len
+ # cu_lengths.append(current_length)
+ # if sample_len > max_length:
+ # max_length = sample_len
+
+ # # If nothing was added, return a minimal packed sample
+ # if not packed_tokens_parts:
+ # return ImageTaskSamplePacked(
+ # __key__="",
+ # __restore_key__=(),
+ # __subflavor__=None,
+ # __subflavors__=samples[0].__subflavors__ if samples else {},
+ # tokens=torch.zeros((0,), dtype=torch.int64),
+ # labels=torch.zeros((0,), dtype=torch.int64),
+ # attn_mask=torch.zeros((0,), dtype=torch.bool),
+ # imgs=packed_imgs,
+ # pixel_values_videos=packed_videos,
+ # cu_lengths=torch.tensor(cu_lengths, dtype=torch.int32),
+ # max_length=0,
+ # num_tiles=[n for s in samples for n in s.num_tiles] if samples else [],
+ # )
+
+ # # For concatenation we need to ensure same dtypes. Replace None slices with zero-length tensors.
+ # def to_tensor_or_empty(x, dtype=None):
+ # if x is None:
+ # return torch.zeros((0,), dtype=dtype if dtype is not None else torch.int64)
+ # if isinstance(x, torch.Tensor):
+ # return x
+ # # try to convert numpy/other to tensor
+ # try:
+ # return torch.as_tensor(x)
+ # except Exception:
+ # return torch.zeros((0,), dtype=dtype if dtype is not None else torch.int64)
+
+ # # infer dtypes from first available slices
+ # first_tokens = next((p for p in packed_tokens_parts if p is not None), None)
+ # first_labels = next((p for p in packed_labels_parts if p is not None), None)
+ # first_masks = next((p for p in packed_masks_parts if p is not None), None)
+
+ # tokens_dtype = first_tokens.dtype if isinstance(first_tokens, torch.Tensor) else torch.int64
+ # labels_dtype = first_labels.dtype if isinstance(first_labels, torch.Tensor) else torch.int64
+ # masks_dtype = first_masks.dtype if isinstance(first_masks, torch.Tensor) else torch.bool
+
+ # packed_tokens = torch.cat([to_tensor_or_empty(p, dtype=tokens_dtype) for p in packed_tokens_parts], dim=0)
+ # packed_labels = torch.cat([to_tensor_or_empty(p, dtype=labels_dtype) for p in packed_labels_parts], dim=0)
+ # packed_masks = torch.cat([to_tensor_or_empty(p, dtype=masks_dtype) for p in packed_masks_parts], dim=0)
+
+ # return ImageTaskSamplePacked(
+ # __key__=",".join([s.__key__ for s in samples]),
+ # __restore_key__=(), # Will be set by energon if needed
+ # __subflavor__=None,
+ # __subflavors__=samples[0].__subflavors__ if samples else {},
+ # tokens=packed_tokens,
+ # labels=packed_labels,
+ # attn_mask=packed_masks,
+ # imgs=packed_imgs,
+ # pixel_values_videos=packed_videos,
+ # cu_lengths=torch.tensor(cu_lengths, dtype=torch.int32),
+ # max_length=max_length,
+ # num_tiles=[n for s in samples for n in s.num_tiles],
+ # )
@stateless
def pack_selected_samples(self, samples: List[ImageTaskSample]) -> List[ImageTaskSamplePacked]:
- """
- Function to pack a list of ImageTaskSample into a single ImageTaskSamplePacked.
-
- NOTE: Energon dataloader calls this method internally if packing is used.
- Please see https://nvidia.github.io/Megatron-Energon/packing.html
-
- Args:
- samples: List of ImageTaskSample instances to pack into one sample.
-
- Returns:
- ImageTaskSamplePacked instance.
- """
-
- packing_seq_len = self.args.seq_length
+ packing_seq_len = int(self.args.seq_length)
packed_tokens = []
packed_labels = []
@@ -336,43 +558,141 @@ def pack_selected_samples(self, samples: List[ImageTaskSample]) -> List[ImageTas
max_length = 0
cu_lengths = [0]
- # Process each sample and build lists that we will concatenate to create the packed sample.
- for _, sample in enumerate(samples):
- sample_len = sample.total_len
-
- if sample_len > max_length:
- max_length = sample_len
-
- # If adding this sample exceeds the max length, stop.
- # This should not happen.
- # The select_samples_to_pack method should have already ensured that the samples fit.
- if current_length + sample_len > packing_seq_len:
- print(f"packing_seq_len:{packing_seq_len}----<<<<<----{current_length + sample_len}")
- raise ValueError(f"Packed sample exceeds the maximum sequence length of {packing_seq_len}: {samples}")
-
- # Add the sample's tokens and labels
- packed_tokens.append(sample.tokens)
- packed_labels.append(sample.labels)
- packed_masks.append(sample.attn_mask)
-
- # Add the images
- if sample.imgs is not None:
- packed_imgs += sample.imgs
- if sample.pixel_values_videos is not None:
- packed_videos += sample.pixel_values_videos
- current_length += sample_len
+ def seq_len_from_tensor(tensor):
+ if tensor is None:
+ return 0
+ if tensor.ndim == 1:
+ return tensor.shape[0]
+ if tensor.ndim >= 2:
+ if tensor.shape[0] == 1 and tensor.shape[1] > 1:
+ return tensor.shape[1]
+ return tensor.shape[0]
+ return tensor.shape[0]
+
+ def slice_by_seq(tensor, keep_len):
+ if tensor is None:
+ return None
+ if keep_len <= 0:
+ if tensor.ndim == 1:
+ return tensor.new_empty((0,))
+ seq_axis = 1 if (tensor.ndim >= 2 and tensor.shape[0] == 1 and tensor.shape[1] > 1) else 0
+ shape = list(tensor.shape)
+ shape[seq_axis] = 0
+ return tensor.new_empty(tuple(shape))
+ if tensor.ndim == 1:
+ return tensor[:keep_len]
+ seq_axis = 1 if (tensor.ndim >= 2 and tensor.shape[0] == 1 and tensor.shape[1] > 1) else 0
+ if seq_axis == 0:
+ return tensor[:keep_len, ...]
+ else:
+ sl = [slice(None)] * tensor.ndim
+ sl[1] = slice(0, keep_len)
+ return tensor[tuple(sl)]
+
+ for idx, sample in enumerate(samples):
+ tokens = getattr(sample, "tokens", None)
+ labels = getattr(sample, "labels", None)
+ attn_mask = getattr(sample, "attn_mask", None)
+
+ candidate_lens = []
+ if tokens is not None:
+ candidate_lens.append(seq_len_from_tensor(tokens))
+ if labels is not None:
+ candidate_lens.append(seq_len_from_tensor(labels))
+ if attn_mask is not None:
+ candidate_lens.append(seq_len_from_tensor(attn_mask))
+ if hasattr(sample, "total_len"):
+ try:
+ candidate_lens.append(int(sample.total_len))
+ except Exception:
+ pass
+
+ if len(candidate_lens) == 0:
+ raise ValueError(f"Sample {idx} has no sequence info (tokens/labels/attn_mask/total_len).")
+
+ orig_len = max(candidate_lens)
+ remaining = packing_seq_len - current_length
+ if remaining <= 0:
+ break
+
+ keep_len = min(orig_len, remaining, packing_seq_len)
+
+ if keep_len < orig_len:
+ _rank0_print(f"[pack] truncating sample idx={idx} from {orig_len} -> {keep_len} (remaining={remaining}) key={getattr(sample, '__key__', 'N/A')})")
+
+ t_tokens = slice_by_seq(tokens, keep_len) if tokens is not None else None
+ t_labels = slice_by_seq(labels, keep_len) if labels is not None else None
+ t_mask = slice_by_seq(attn_mask, keep_len) if attn_mask is not None else None
+
+ # --- FIXED: choose first non-None tensor explicitly, don't use `or` with tensors ---
+ first_available = None
+ for cand in (t_tokens, t_labels, t_mask):
+ if cand is not None:
+ first_available = cand
+ break
+
+ tlen = 0 if first_available is None else seq_len_from_tensor(first_available)
+
+ # Ensure tlen <= keep_len (safety)
+ if tlen > keep_len:
+ # Force-slice to keep_len
+ if t_tokens is not None:
+ t_tokens = slice_by_seq(t_tokens, keep_len)
+ if t_labels is not None:
+ t_labels = slice_by_seq(t_labels, keep_len)
+ if t_mask is not None:
+ t_mask = slice_by_seq(t_mask, keep_len)
+ tlen = keep_len
+
+ if tlen > max_length:
+ max_length = tlen
+
+ if t_tokens is None:
+ raise ValueError(f"Sample {idx} missing tokens after slicing (key={getattr(sample,'__key__','N/A')}).")
+ if t_labels is None:
+ raise ValueError(f"Sample {idx} missing labels after slicing (key={getattr(sample,'__key__','N/A')}).")
+ if t_mask is None:
+ raise ValueError(f"Sample {idx} missing attn_mask after slicing (key={getattr(sample,'__key__','N/A')}).")
+
+ packed_tokens.append(t_tokens)
+ packed_labels.append(t_labels)
+ packed_masks.append(t_mask)
+
+ if getattr(sample, "imgs", None):
+ packed_imgs += list(sample.imgs)
+ if getattr(sample, "pixel_values_videos", None):
+ packed_videos += list(sample.pixel_values_videos)
+
+ current_length += tlen
cu_lengths.append(current_length)
- # Concatenate packed tokens and labels.
- packed_tokens = torch.cat(packed_tokens, dim=0)
- packed_labels = torch.cat(packed_labels, dim=0)
- packed_masks = torch.cat(packed_masks, dim=0)
+ if current_length >= packing_seq_len:
+ break
+
+ if len(packed_tokens) == 0:
+ packed_tokens = torch.empty((0,), dtype=torch.long)
+ packed_labels = torch.empty((0,), dtype=torch.long)
+ packed_masks = torch.empty((0,), dtype=torch.bool)
+ else:
+ packed_tokens = torch.cat(packed_tokens, dim=0)
+ packed_labels = torch.cat(packed_labels, dim=0)
+ packed_masks = torch.cat(packed_masks, dim=0)
+
+ final_len = packed_tokens.shape[0] if packed_tokens is not None else 0
+ if final_len > packing_seq_len:
+ packed_tokens = packed_tokens[:packing_seq_len]
+ packed_labels = packed_labels[:packing_seq_len]
+ packed_masks = packed_masks[:packing_seq_len]
+ final_len = packing_seq_len
+ # best-effort adjust cu_lengths (keep it simple)
+ if cu_lengths:
+ cu_lengths[-1] = final_len
return ImageTaskSamplePacked(
__key__=",".join([s.__key__ for s in samples]),
- __restore_key__=(), # Will be set by energon based on `samples`
+ __restore_key__=(),
__subflavor__=None,
- __subflavors__=samples[0].__subflavors__,
+ __subflavors__=samples[0].__subflavors__ if samples else {},
tokens=packed_tokens,
labels=packed_labels,
attn_mask=packed_masks,
@@ -383,7 +703,6 @@ def pack_selected_samples(self, samples: List[ImageTaskSample]) -> List[ImageTas
num_tiles=[n for s in samples for n in s.num_tiles],
)
-
def print_error_handler(exc: Exception, key: Optional[str]):
""" Print error handler function called when an exception occurs during loading. """
print(
diff --git a/aiak_training_llm/models/custom/common/local_norm.py b/aiak_training_llm/models/custom/common/local_norm.py
index d488ee24..f284538b 100644
--- a/aiak_training_llm/models/custom/common/local_norm.py
+++ b/aiak_training_llm/models/custom/common/local_norm.py
@@ -1,75 +1,149 @@
-"""megatron local norm"""
+"""megatron local norm - Apex NOT used; pure-PyTorch fallbacks."""
+from typing import Optional
+import torch
+import math
+import torch.nn as nn
+from torch.nn.parameter import Parameter
+
+# Megatron config type (import only for typing/runtime usage)
from megatron.core.transformer.transformer_config import TransformerConfig
-from megatron.core.fusions.fused_layer_norm import FusedLayerNorm
+# Prefer Megatron's fused op if present (this is NOT Apex):
try:
- from apex.normalization.fused_layer_norm import FusedRMSNorm as ApexFusedRMSNorm
- HAVE_FUSED_RMS_NORM = True
-except:
- HAVE_FUSED_RMS_NORM = False
+ from megatron.core.fusions.fused_layer_norm import FusedLayerNorm # optional
+ HAVE_MEGATRON_FUSED_LAYER_NORM = True
+except Exception:
+ HAVE_MEGATRON_FUSED_LAYER_NORM = False
-try:
- from apex.normalization.fused_layer_norm import FusedLayerNorm as ApexFusedLayerNorm
- HAVE_FUSED_LAYER_NORM = True
-except:
- HAVE_FUSED_LAYER_NORM = False
+# --- Pure PyTorch RMSNorm implementation ---
+class RMSNorm(nn.Module):
+ """Pure-PyTorch RMSNorm implementation.
+ y = x * (weight / sqrt(mean(x^2, dim=-1, keepdim=True) + eps))
+
+ Keeps a `weight` attribute compatible with other code expecting per-channel
+ affine parameters. Sets `weight.sequence_parallel` when a config is provided.
+ """
-class FusedRMSNorm(ApexFusedRMSNorm):
- """Fused RMS Norm"""
- def __init__(self,
- config: TransformerConfig,
- hidden_size: int,
- eps=1e-5,
- elementwise_affine=True):
+ def __init__(
+ self,
+ config: TransformerConfig,
+ hidden_size: int,
+ eps: float = 1e-5,
+ elementwise_affine: bool = True,
+ ):
+ super().__init__()
+ self.config = config
+ self.hidden_size = hidden_size
+ self.eps = eps
+ self.elementwise_affine = elementwise_affine
- if not HAVE_FUSED_RMS_NORM:
- # TODO: Add pytorch only rms norm
- raise ValueError(f'Apex must currently be installed to use FusedRMSNorm op.')
+ if self.elementwise_affine:
+ self.weight = Parameter(torch.ones(hidden_size))
+ else:
+ # keep attribute for compatibility
+ self.register_buffer("weight", torch.ones(hidden_size))
- super().__init__(hidden_size,
- eps=eps,
- elementwise_affine=elementwise_affine)
+ # optional support for megatron sequence_parallel flag on config
+ self.sequence_parallel = getattr(self.config, "sequence_parallel", False)
+ try:
+ setattr(self.weight, "sequence_parallel", self.sequence_parallel)
+ except Exception:
+ # some frameworks may not allow setting attributes on Parameter on certain dtypes,
+ # ignore silently (compatibility only)
+ pass
- self.config = config
+ def forward(self, x: torch.Tensor):
+ # compute RMS over last dim
+ # ms shape: (..., 1)
+ ms = x.pow(2).mean(dim=-1, keepdim=True)
+ rs = torch.rsqrt(ms + self.eps) # (..., 1)
+ if self.elementwise_affine:
+ # broadcast weight across leading dims; if hidden size differs from weight length,
+ # repeat or truncate to fit to avoid shape mismatch.
+ hidden = x.size(-1)
+ w_raw = self.weight
+ if w_raw.numel() != hidden:
+ repeat = math.ceil(hidden / w_raw.numel())
+ w_raw = w_raw.repeat(repeat)[:hidden]
+ w = w_raw.view(*([1] * (x.dim() - 1)), -1)
+ # return x * rs * w
+ return x * rs * w
+ else:
+ return x * rs
- self.sequence_parallel = self.config.sequence_parallel
- # set sequence parallelism flag on weight parameters
- setattr(self.weight, 'sequence_parallel', self.sequence_parallel)
+# --- LayerNorm wrapper (Megatron fused if available, else torch.nn.LayerNorm) ---
+class TorchLayerNormWrapper(nn.Module):
+ """Wrap torch.nn.LayerNorm so its initializer and attributes match expected API."""
+ def __init__(self, hidden_size: int, eps: float = 1e-5, elementwise_affine: bool = True):
+ super().__init__()
+ self.layernorm = nn.LayerNorm(hidden_size, eps=eps, elementwise_affine=elementwise_affine)
+ if elementwise_affine:
+ # expose weight attribute for compatibility (so callers can set sequence_parallel)
+ self.weight = self.layernorm.weight
+ else:
+ # still expose a weight-like buffer for code that expects it
+ self.register_buffer("weight", torch.ones(hidden_size))
+
+ def forward(self, x: torch.Tensor):
+ return self.layernorm(x)
+
+
+# Pick the best LayerNorm implementation available (but do NOT import Apex anywhere).
+if HAVE_MEGATRON_FUSED_LAYER_NORM:
+ LayerNormClass = FusedLayerNorm
+else:
+ LayerNormClass = TorchLayerNormWrapper
+
+# --- Factory class: LocalNorm ---
class LocalNorm:
"""
- A conditional wrapper to initialize an instance of Megatron Local `LayerNorm` or `RMSNorm` based on input
+ Factory to create a normalization module consistent with megatron config.
+
+ Usage:
+ norm = LocalNorm(config, hidden_size, eps=1e-5, elementwise_affine=True)
+
+ This will return either:
+ - a LayerNorm-like module (Megatron fused if available, otherwise torch.nn.LayerNorm wrapper), or
+ - RMSNorm (pure-PyTorch).
"""
- # TODO should we ditch normalization config and just use spec to choose LayerNorm vs RMSNorm?
- def __new__(cls, config: TransformerConfig, hidden_size: int, eps: float = 1e-5, elementwise_affine=True):
- if config.normalization == "LayerNorm":
- if elementwise_affine:
- instance = FusedLayerNorm(
- config=config,
- hidden_size=hidden_size,
- eps=eps,
- )
+ def __new__(
+ cls,
+ config: TransformerConfig,
+ hidden_size: int,
+ eps: float = 1e-5,
+ elementwise_affine: bool = True,
+ ):
+ norm_type = getattr(config, "normalization", "LayerNorm")
+ if norm_type == "LayerNorm":
+ # prefer megatron fused if it exists, else torch wrapper
+ # Megatron FusedLayerNorm initializer signature: (config=config, hidden_size=..., eps=...)
+ if HAVE_MEGATRON_FUSED_LAYER_NORM:
+ # Create using Megatron fused implementation (keeps config semantics)
+ try:
+ instance = FusedLayerNorm(config=config, hidden_size=hidden_size, eps=eps)
+ except Exception:
+ # fallback to torch wrapper if megatron fused can't be constructed
+ instance = LayerNormClass(hidden_size, eps=eps, elementwise_affine=elementwise_affine)
else:
- assert HAVE_FUSED_LAYER_NORM, "Apex must currently be installed to use FusedLayerNorm op."
- instance = ApexFusedLayerNorm(
- hidden_size,
- eps=eps,
- elementwise_affine=elementwise_affine,
- )
- elif config.normalization == "RMSNorm":
- instance = FusedRMSNorm(
- config=config,
- hidden_size=hidden_size,
- eps=eps,
- elementwise_affine=elementwise_affine,
- )
+ instance = LayerNormClass(hidden_size, eps=eps, elementwise_affine=elementwise_affine)
+
+ # if config exposes sequence_parallel, replicate attribute on weight if present
+ seq_par = getattr(config, "sequence_parallel", False)
+ if hasattr(instance, "weight"):
+ try:
+ setattr(instance.weight, "sequence_parallel", seq_par)
+ except Exception:
+ pass
+ elif norm_type == "RMSNorm":
+ instance = RMSNorm(config=config, hidden_size=hidden_size, eps=eps, elementwise_affine=elementwise_affine)
else:
- raise Exception('Only LayerNorm and RMSNorm are curently supported')
+ raise ValueError("Only 'LayerNorm' and 'RMSNorm' are supported for LocalNorm")
return instance
diff --git a/aiak_training_llm/models/custom/common/local_norm_org.py b/aiak_training_llm/models/custom/common/local_norm_org.py
new file mode 100644
index 00000000..cbf7e8fc
--- /dev/null
+++ b/aiak_training_llm/models/custom/common/local_norm_org.py
@@ -0,0 +1,77 @@
+"""megatron local norm"""
+
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.fusions.fused_layer_norm import FusedLayerNorm
+
+try:
+ from apex.normalization.fused_layer_norm import FusedRMSNorm as ApexFusedRMSNorm
+ HAVE_FUSED_RMS_NORM = True
+except Exception as e:
+ HAVE_FUSED_RMS_NORM = False
+ print("Failed to import Apex FusedRMSNorm:", repr(e))
+print("HAVE_FUSED_RMS_NORM =", HAVE_FUSED_RMS_NORM)
+try:
+ from apex.normalization.fused_layer_norm import FusedLayerNorm as ApexFusedLayerNorm
+ HAVE_FUSED_LAYER_NORM = True
+except:
+ HAVE_FUSED_LAYER_NORM = False
+print("HAVE_FUSED_LAYER_NORM =", HAVE_FUSED_LAYER_NORM)
+
+class FusedRMSNorm(ApexFusedRMSNorm):
+ """Fused RMS Norm"""
+ def __init__(self,
+ config: TransformerConfig,
+ hidden_size: int,
+ eps=1e-5,
+ elementwise_affine=True):
+
+ if not HAVE_FUSED_RMS_NORM:
+ # TODO: Add pytorch only rms norm
+ raise ValueError(f'Apex must currently be installed to use FusedRMSNorm op.')
+
+ super().__init__(hidden_size,
+ eps=eps,
+ elementwise_affine=elementwise_affine)
+
+ self.config = config
+
+ self.sequence_parallel = self.config.sequence_parallel
+
+ # set sequence parallelism flag on weight parameters
+ setattr(self.weight, 'sequence_parallel', self.sequence_parallel)
+
+
+class LocalNorm:
+ """
+ A conditional wrapper to initialize an instance of Megatron Local `LayerNorm` or `RMSNorm` based on input
+ """
+
+ # TODO should we ditch normalization config and just use spec to choose LayerNorm vs RMSNorm?
+ def __new__(cls, config: TransformerConfig, hidden_size: int, eps: float = 1e-5, elementwise_affine=True):
+ if config.normalization == "LayerNorm":
+ if elementwise_affine:
+ print("LayerNorm")
+ instance = FusedLayerNorm(
+ config=config,
+ hidden_size=hidden_size,
+ eps=eps,
+ )
+ else:
+ assert HAVE_FUSED_LAYER_NORM, "Apex must currently be installed to use FusedLayerNorm op."
+ instance = ApexFusedLayerNorm(
+ hidden_size,
+ eps=eps,
+ elementwise_affine=elementwise_affine,
+ )
+ elif config.normalization == "RMSNorm":
+ instance = FusedRMSNorm(
+ config=config,
+ hidden_size=hidden_size,
+ eps=eps,
+ elementwise_affine=elementwise_affine,
+ )
+
+ else:
+ raise Exception('Only LayerNorm and RMSNorm are curently supported')
+
+ return instance
diff --git a/aiak_training_llm/models/dispatch.py b/aiak_training_llm/models/dispatch.py
index ac4fc063..54cdb7a6 100644
--- a/aiak_training_llm/models/dispatch.py
+++ b/aiak_training_llm/models/dispatch.py
@@ -49,6 +49,7 @@ def _gpu_backend_transformer_layer_modules() -> MultiAccModules:
TENorm,
TELinear,
)
+ from megatron.core.transformer.dot_product_attention import DotProductAttention
from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add
from megatron.core.models.common.embeddings.rotary_pos_embedding import apply_rotary_pos_emb
@@ -56,6 +57,11 @@ def _gpu_backend_transformer_layer_modules() -> MultiAccModules:
args = get_args()
+ # Use local DotProductAttention if transformer_impl is "local", otherwise use TE version
+ # DotProductAttention=TEDotProductAttention (old hardcoded version, now commented)
+ use_te_attn = args.transformer_impl != "local"
+ attn_module = TEDotProductAttention if use_te_attn else DotProductAttention
+
return MultiAccModules(
# dense linear
TELayerNormColumnParallelLinear=TELayerNormColumnParallelLinear,
@@ -68,7 +74,7 @@ def _gpu_backend_transformer_layer_modules() -> MultiAccModules:
ColumnParallelLinear=ColumnParallelLinear,
RowParallelLinear=RowParallelLinear,
# attn
- DotProductAttention=TEDotProductAttention,
+ DotProductAttention=attn_module,
# norm
TENorm=TENorm,
LocalNorm=LocalNorm,
diff --git a/aiak_training_llm/models/fastvit/__init__.py b/aiak_training_llm/models/fastvit/__init__.py
new file mode 100644
index 00000000..04e6a56c
--- /dev/null
+++ b/aiak_training_llm/models/fastvit/__init__.py
@@ -0,0 +1,10 @@
+#
+# For licensing see accompanying LICENSE file.
+# Copyright (C) 2025 Apple Inc. All Rights Reserved.
+#
+
+from . import mobileclip
+from .mobileclip_encoder import MobileCLIPVisionTower
+from .fastvit_vision_model import FastViTModel
+
+__all__ = ['mobileclip', 'MobileCLIPVisionTower', 'FastViTModel']
diff --git a/aiak_training_llm/models/fastvit/builder.py b/aiak_training_llm/models/fastvit/builder.py
new file mode 100644
index 00000000..290ecdf7
--- /dev/null
+++ b/aiak_training_llm/models/fastvit/builder.py
@@ -0,0 +1,19 @@
+import os
+from .clip_encoder import CLIPVisionTower, CLIPVisionTowerS2
+from .mobileclip_encoder import MobileCLIPVisionTower
+
+
+def build_vision_tower(vision_tower_cfg, **kwargs):
+ vision_tower = getattr(vision_tower_cfg, 'mm_vision_tower', getattr(vision_tower_cfg, 'vision_tower', None))
+ is_absolute_path_exists = os.path.exists(vision_tower)
+ use_s2 = getattr(vision_tower_cfg, 's2', False)
+
+ if is_absolute_path_exists or vision_tower.startswith("openai") or vision_tower.startswith("laion") or "ShareGPT4V" in vision_tower:
+ if use_s2:
+ return CLIPVisionTowerS2(vision_tower, args=vision_tower_cfg, **kwargs)
+ else:
+ return CLIPVisionTower(vision_tower, args=vision_tower_cfg, **kwargs)
+ elif "mobileclip" in vision_tower.lower():
+ return MobileCLIPVisionTower(vision_tower, args=vision_tower_cfg, **kwargs)
+
+ raise ValueError(f'Unknown vision tower: {vision_tower}')
\ No newline at end of file
diff --git a/aiak_training_llm/models/fastvit/constants.py b/aiak_training_llm/models/fastvit/constants.py
new file mode 100644
index 00000000..6964799b
--- /dev/null
+++ b/aiak_training_llm/models/fastvit/constants.py
@@ -0,0 +1,13 @@
+CONTROLLER_HEART_BEAT_EXPIRATION = 30
+WORKER_HEART_BEAT_INTERVAL = 15
+
+LOGDIR = "."
+
+# Model Constants
+IGNORE_INDEX = -100
+IMAGE_TOKEN_INDEX = -200
+DEFAULT_IMAGE_TOKEN = ""
+DEFAULT_IMAGE_PATCH_TOKEN = ""
+DEFAULT_IM_START_TOKEN = ""
+DEFAULT_IM_END_TOKEN = ""
+IMAGE_PLACEHOLDER = ""
\ No newline at end of file
diff --git a/aiak_training_llm/models/fastvit/fastvit_preprocessor.py b/aiak_training_llm/models/fastvit/fastvit_preprocessor.py
new file mode 100644
index 00000000..5cf47d71
--- /dev/null
+++ b/aiak_training_llm/models/fastvit/fastvit_preprocessor.py
@@ -0,0 +1,69 @@
+#
+# FastViT Image Preprocessing
+# Following Apple FastVLM preprocessing approach
+#
+
+from PIL import Image
+import torch
+from transformers import CLIPImageProcessor
+
+
+class FastViTImageProcessor:
+ """
+ Image processor for FastViT following FastVLM's approach.
+ Uses CLIPImageProcessor with specific settings for FastViT.
+ """
+
+ def __init__(self, image_size=1024):
+ """
+ Initialize FastViT image processor.
+
+ Args:
+ image_size: Input image size (default 384 as in FastVLM)
+ """
+ self.image_size = image_size
+
+ # Create CLIPImageProcessor with FastViT settings
+ # Following mobileclip_encoder.py from FastVLM
+ # Source: ml-fastvlm/llava/model/multimodal_encoder/mobileclip_encoder.py (lines 45-49)
+ self.image_mean = [0.0, 0.0, 0.0] # No normalization (black padding for 'pad' mode)
+ self.image_std = [1.0, 1.0, 1.0]
+
+ self.processor = CLIPImageProcessor(
+ crop_size={"height": image_size, "width": image_size},
+ image_mean=self.image_mean,
+ image_std=self.image_std,
+ size={"shortest_edge": image_size}
+ )
+
+ def preprocess(self, image):
+ """
+ Preprocess a single image for FastViT.
+
+ Args:
+ image: PIL Image
+
+ Returns:
+ Preprocessed image tensor
+ """
+ if isinstance(image, Image.Image):
+ # Process with CLIPImageProcessor
+ processed = self.processor(images=image, return_tensors="pt")
+ return processed['pixel_values'][0] # Remove batch dimension
+ else:
+ raise ValueError(f"Expected PIL Image, got {type(image)}")
+
+ def __call__(self, images):
+ """
+ Process single image or batch of images.
+
+ Args:
+ images: PIL Image or list of PIL Images
+
+ Returns:
+ Tensor of preprocessed images
+ """
+ if isinstance(images, list):
+ return torch.stack([self.preprocess(img) for img in images])
+ else:
+ return self.preprocess(images).unsqueeze(0)
diff --git a/aiak_training_llm/models/fastvit/fastvit_vision_model.py b/aiak_training_llm/models/fastvit/fastvit_vision_model.py
new file mode 100644
index 00000000..85f63ee2
--- /dev/null
+++ b/aiak_training_llm/models/fastvit/fastvit_vision_model.py
@@ -0,0 +1,102 @@
+#
+# FastViT Vision Model wrapper for Megatron integration
+# Adapted from Apple's FastVLM
+#
+
+import torch
+import torch.nn as nn
+from megatron.core.transformer import MegatronModule
+
+from aiak_training_llm.models.fastvit.mobileclip_encoder import MobileCLIPVisionTower
+
+
+class FastViTModel(MegatronModule):
+ """
+ FastViT Vision Model wrapper for Megatron.
+
+ This wraps the MobileCLIPVisionTower (which contains FastViTHD)
+ to be compatible with Megatron's training framework.
+
+ Args:
+ config: Megatron TransformerConfig
+ layer_spec: Module specification (not used for FastViT)
+ """
+
+ def __init__(self, config, layer_spec):
+ # Initialize MegatronModule (will set self.config)
+ super().__init__(config=config)
+
+ # Create a simple args object for MobileCLIPVisionTower
+ class Args:
+ def __init__(self, unfreeze_mm_vision_tower):
+ self.unfreeze_mm_vision_tower = unfreeze_mm_vision_tower
+
+ args = Args(getattr(config, 'unfreeze_mm_vision_tower', False))
+
+ # FastViTHD model name with resolution
+ # Format: mobileclip_l_{resolution}
+ # The first two parts must match the config file name (mobileclip_l.json)
+ # The resolution is extracted from the last part
+ vision_tower_name = "mobileclip_l_1024"
+
+ # Override with config if provided
+ if hasattr(config, 'vision_tower_name'):
+ vision_tower_name = config.vision_tower_name
+ elif hasattr(config, 'img_h') and config.img_w == config.img_h:
+ # Use square image size from config
+ vision_tower_name = f"mobileclip_l_{config.img_h}"
+
+ # Initialize the FastViT vision tower
+ self.vision_tower = MobileCLIPVisionTower(
+ vision_tower=vision_tower_name,
+ args=args,
+ delay_load=False
+ )
+
+ self._hidden_size = self.vision_tower.hidden_size
+
+ def forward(self, images, grid_thw=None):
+ """
+ Forward pass through FastViT.
+
+ Args:
+ images: Image tensor [batch, channels, height, width]
+ grid_thw: Grid dimensions (not used by FastViT, kept for API compatibility)
+
+ Returns:
+ Vision features [total_image_tokens, hidden_size]
+ window_index: None (FastViT doesn't use windowing)
+ """
+ # MobileCLIPVisionTower expects single images or list of images
+ # For batch processing, we need to handle it appropriately
+ if images.dim() == 4:
+ # Batch of images [B, C, H, W]
+ image_features = self.vision_tower.forward_images(images)
+ else:
+ raise ValueError(f"Unexpected image tensor shape: {images.shape}")
+
+ if image_features.dim() == 3:
+ # [num_images, tokens_per_image, hidden] -> [all_image_tokens, hidden]
+ image_features = image_features.flatten(0, 1).contiguous()
+
+ # Return features and None for window_index (compatibility with Qwen2-VL API)
+ return image_features, None
+
+ def set_input_tensor(self, input_tensor):
+ """
+ Set input tensor.
+
+ FastViT doesn't use pipeline parallelism, so this is a no-op.
+ Required for compatibility with Megatron's pipeline parallel interface.
+ """
+ pass
+
+ @property
+ def hidden_size(self):
+ """Return the hidden size of the vision model output."""
+ return self._hidden_size
+
+ @property
+ def vision_config(self):
+ """Return vision tower config."""
+ return self.vision_tower.config
diff --git a/aiak_training_llm/models/fastvit/mm_utils.py b/aiak_training_llm/models/fastvit/mm_utils.py
new file mode 100644
index 00000000..d08dd58d
--- /dev/null
+++ b/aiak_training_llm/models/fastvit/mm_utils.py
@@ -0,0 +1,250 @@
+import PIL
+from PIL import Image
+PIL.Image.MAX_IMAGE_PIXELS=500000000
+from io import BytesIO
+import base64
+import torch
+import math
+import ast
+
+from transformers import StoppingCriteria
+from aiak_training_llm.models.fastvit.constants import IMAGE_TOKEN_INDEX
+
+
+def select_best_resolution(original_size, possible_resolutions):
+ """
+ Selects the best resolution from a list of possible resolutions based on the original size.
+
+ Args:
+ original_size (tuple): The original size of the image in the format (width, height).
+ possible_resolutions (list): A list of possible resolutions in the format [(width1, height1), (width2, height2), ...].
+
+ Returns:
+ tuple: The best fit resolution in the format (width, height).
+ """
+ original_width, original_height = original_size
+ best_fit = None
+ max_effective_resolution = 0
+ min_wasted_resolution = float('inf')
+
+ for width, height in possible_resolutions:
+ scale = min(width / original_width, height / original_height)
+ downscaled_width, downscaled_height = int(original_width * scale), int(original_height * scale)
+ effective_resolution = min(downscaled_width * downscaled_height, original_width * original_height)
+ wasted_resolution = (width * height) - effective_resolution
+
+ if effective_resolution > max_effective_resolution or (effective_resolution == max_effective_resolution and wasted_resolution < min_wasted_resolution):
+ max_effective_resolution = effective_resolution
+ min_wasted_resolution = wasted_resolution
+ best_fit = (width, height)
+
+ return best_fit
+
+
+def resize_and_pad_image(image, target_resolution):
+ """
+ Resize and pad an image to a target resolution while maintaining aspect ratio.
+
+ Args:
+ image (PIL.Image.Image): The input image.
+ target_resolution (tuple): The target resolution (width, height) of the image.
+
+ Returns:
+ PIL.Image.Image: The resized and padded image.
+ """
+ original_width, original_height = image.size
+ target_width, target_height = target_resolution
+
+ scale_w = target_width / original_width
+ scale_h = target_height / original_height
+
+ if scale_w < scale_h:
+ new_width = target_width
+ new_height = min(math.ceil(original_height * scale_w), target_height)
+ else:
+ new_height = target_height
+ new_width = min(math.ceil(original_width * scale_h), target_width)
+
+ # Resize the image
+ resized_image = image.resize((new_width, new_height))
+
+ new_image = Image.new('RGB', (target_width, target_height), (0, 0, 0))
+ paste_x = (target_width - new_width) // 2
+ paste_y = (target_height - new_height) // 2
+ new_image.paste(resized_image, (paste_x, paste_y))
+
+ return new_image
+
+
+def divide_to_patches(image, patch_size):
+ """
+ Divides an image into patches of a specified size.
+
+ Args:
+ image (PIL.Image.Image): The input image.
+ patch_size (int): The size of each patch.
+
+ Returns:
+ list: A list of PIL.Image.Image objects representing the patches.
+ """
+ patches = []
+ width, height = image.size
+ for i in range(0, height, patch_size):
+ for j in range(0, width, patch_size):
+ box = (j, i, j + patch_size, i + patch_size)
+ patch = image.crop(box)
+ patches.append(patch)
+
+ return patches
+
+
+def get_anyres_image_grid_shape(image_size, grid_pinpoints, patch_size):
+ """
+ Calculate the shape of the image patch grid after the preprocessing for images of any resolution.
+
+ Args:
+ image_size (tuple): The size of the input image in the format (width, height).
+ grid_pinpoints (str): A string representation of a list of possible resolutions.
+ patch_size (int): The size of each image patch.
+
+ Returns:
+ tuple: The shape of the image patch grid in the format (width, height).
+ """
+ if type(grid_pinpoints) is list:
+ possible_resolutions = grid_pinpoints
+ else:
+ possible_resolutions = ast.literal_eval(grid_pinpoints)
+ width, height = select_best_resolution(image_size, possible_resolutions)
+ return width // patch_size, height // patch_size
+
+
+def process_anyres_image(image, processor, grid_pinpoints):
+ """
+ Process an image with variable resolutions.
+
+ Args:
+ image (PIL.Image.Image): The input image to be processed.
+ processor: The image processor object.
+ grid_pinpoints (str): A string representation of a list of possible resolutions.
+
+ Returns:
+ torch.Tensor: A tensor containing the processed image patches.
+ """
+ if type(grid_pinpoints) is list:
+ possible_resolutions = grid_pinpoints
+ else:
+ possible_resolutions = ast.literal_eval(grid_pinpoints)
+ best_resolution = select_best_resolution(image.size, possible_resolutions)
+ image_padded = resize_and_pad_image(image, best_resolution)
+
+ patches = divide_to_patches(image_padded, processor.crop_size['height'])
+
+ image_original_resize = image.resize((processor.size['shortest_edge'], processor.size['shortest_edge']))
+
+ image_patches = [image_original_resize] + patches
+ image_patches = [processor.preprocess(image_patch, return_tensors='pt')['pixel_values'][0]
+ for image_patch in image_patches]
+ return torch.stack(image_patches, dim=0)
+
+
+def load_image_from_base64(image):
+ return Image.open(BytesIO(base64.b64decode(image)))
+
+
+def expand2square(pil_img, background_color):
+ width, height = pil_img.size
+ if width == height:
+ return pil_img
+ elif width > height:
+ result = Image.new(pil_img.mode, (width, width), background_color)
+ result.paste(pil_img, (0, (width - height) // 2))
+ return result
+ else:
+ result = Image.new(pil_img.mode, (height, height), background_color)
+ result.paste(pil_img, ((height - width) // 2, 0))
+ return result
+
+
+def process_images(images, image_processor, model_cfg):
+ image_aspect_ratio = getattr(model_cfg, "image_aspect_ratio", None)
+ new_images = []
+ if image_aspect_ratio == 'pad':
+ for image in images:
+ image = expand2square(image, tuple(int(x*255) for x in image_processor.image_mean))
+ image = image_processor.preprocess(image, return_tensors='pt')['pixel_values'][0]
+ new_images.append(image)
+ elif image_aspect_ratio == "anyres":
+ for image in images:
+ image = process_anyres_image(image, image_processor, model_cfg.image_grid_pinpoints)
+ new_images.append(image)
+ else:
+ return image_processor(images, return_tensors='pt')['pixel_values']
+ if all(x.shape == new_images[0].shape for x in new_images):
+ new_images = torch.stack(new_images, dim=0)
+ return new_images
+
+
+def tokenizer_image_token(prompt, tokenizer, image_token_index=IMAGE_TOKEN_INDEX, return_tensors=None):
+ prompt_chunks = [tokenizer(chunk).input_ids for chunk in prompt.split('')]
+
+ def insert_separator(X, sep):
+ return [ele for sublist in zip(X, [sep]*len(X)) for ele in sublist][:-1]
+
+ input_ids = []
+ offset = 0
+ if len(prompt_chunks) > 0 and len(prompt_chunks[0]) > 0 and prompt_chunks[0][0] == tokenizer.bos_token_id:
+ offset = 1
+ input_ids.append(prompt_chunks[0][0])
+
+ for x in insert_separator(prompt_chunks, [image_token_index] * (offset + 1)):
+ input_ids.extend(x[offset:])
+
+ if return_tensors is not None:
+ if return_tensors == 'pt':
+ return torch.tensor(input_ids, dtype=torch.long)
+ raise ValueError(f'Unsupported tensor type: {return_tensors}')
+ return input_ids
+
+
+def get_model_name_from_path(model_path):
+ model_path = model_path.strip("/")
+ model_paths = model_path.split("/")
+ if model_paths[-1].startswith('checkpoint-'):
+ return model_paths[-2] + "_" + model_paths[-1]
+ else:
+ return model_paths[-1]
+
+
+class KeywordsStoppingCriteria(StoppingCriteria):
+ def __init__(self, keywords, tokenizer, input_ids):
+ self.keywords = keywords
+ self.keyword_ids = []
+ self.max_keyword_len = 0
+ for keyword in keywords:
+ cur_keyword_ids = tokenizer(keyword).input_ids
+ if len(cur_keyword_ids) > 1 and cur_keyword_ids[0] == tokenizer.bos_token_id:
+ cur_keyword_ids = cur_keyword_ids[1:]
+ if len(cur_keyword_ids) > self.max_keyword_len:
+ self.max_keyword_len = len(cur_keyword_ids)
+ self.keyword_ids.append(torch.tensor(cur_keyword_ids))
+ self.tokenizer = tokenizer
+ self.start_len = input_ids.shape[1]
+
+ def call_for_batch(self, output_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+ offset = min(output_ids.shape[1] - self.start_len, self.max_keyword_len)
+ self.keyword_ids = [keyword_id.to(output_ids.device) for keyword_id in self.keyword_ids]
+ for keyword_id in self.keyword_ids:
+ truncated_output_ids = output_ids[0, -keyword_id.shape[0]:]
+ if torch.equal(truncated_output_ids, keyword_id):
+ return True
+ outputs = self.tokenizer.batch_decode(output_ids[:, -offset:], skip_special_tokens=True)[0]
+ for keyword in self.keywords:
+ if keyword in outputs:
+ return True
+ return False
+
+ def __call__(self, output_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+ outputs = []
+ for i in range(output_ids.shape[0]):
+ outputs.append(self.call_for_batch(output_ids[i].unsqueeze(0), scores))
+ return all(outputs)
\ No newline at end of file
diff --git a/aiak_training_llm/models/fastvit/mobileclip/__init__.py b/aiak_training_llm/models/fastvit/mobileclip/__init__.py
new file mode 100644
index 00000000..28e2162c
--- /dev/null
+++ b/aiak_training_llm/models/fastvit/mobileclip/__init__.py
@@ -0,0 +1,86 @@
+#
+# For licensing see accompanying LICENSE file.
+# Copyright (C) 2025 Apple Inc. All Rights Reserved.
+#
+import os
+import json
+from typing import Any
+
+import torch.nn as nn
+from timm.models import create_model
+
+from .mci import GlobalPool2D
+
+
+def load_model_config(
+ model_name: str,
+) -> Any:
+ # Strip suffixes to model name
+ model_name = "_".join(model_name.split("_")[0:2])
+
+ # Config files
+ root_dir = os.path.dirname(os.path.abspath(__file__))
+ configs_dir = os.path.join(root_dir, "configs")
+ model_cfg_file = os.path.join(configs_dir, model_name + ".json")
+
+ # Get config from yaml file
+ if not os.path.exists(model_cfg_file):
+ raise ValueError(f"Unsupported model name: {model_name}")
+ model_cfg = json.load(open(model_cfg_file, "r"))
+
+ return model_cfg
+
+
+class MCi(nn.Module):
+ """
+ This class implements `MCi Models `_
+ """
+
+ def __init__(self, model_name: str, *args, **kwargs) -> None:
+ super().__init__()
+ self.projection_dim = None
+ if "projection_dim" in kwargs:
+ self.projection_dim = kwargs.get("projection_dim")
+
+ # Create model
+ self.model = create_model(model_name, projection_dim=self.projection_dim)
+ # Build out projection head.
+ if self.projection_dim is not None:
+ if hasattr(self.model, "head"):
+ self.model.head = MCi._update_image_classifier(
+ image_classifier=self.model.head, projection_dim=self.projection_dim
+ )
+
+ def forward(self, x: Any, *args, **kwargs) -> Any:
+ """A forward function of the model."""
+ x = self.model(x, *args, **kwargs)
+ return x
+
+ @staticmethod
+ def _get_in_feature_dimension(image_classifier: nn.Module) -> int:
+ """Return the input feature dimension to the image classification head."""
+ in_features = None
+ if isinstance(image_classifier, nn.Sequential):
+ # Classifier that uses nn.Sequential usually has global pooling and
+ # multiple linear layers. Find the first linear layer and get its
+ # in_features
+ for layer in image_classifier:
+ if isinstance(layer, nn.Linear):
+ in_features = layer.in_features
+ break
+ elif isinstance(image_classifier, nn.Linear):
+ in_features = image_classifier.in_features
+
+ if in_features is None:
+ raise NotImplementedError(
+ f"Cannot get input feature dimension of {image_classifier}."
+ )
+ return in_features
+
+ @staticmethod
+ def _update_image_classifier(
+ image_classifier: nn.Module, projection_dim: int, *args, **kwargs
+ ) -> nn.Module:
+ in_features = MCi._get_in_feature_dimension(image_classifier)
+ new_img_classifier = GlobalPool2D(in_dim=in_features, out_dim=projection_dim)
+ return new_img_classifier
\ No newline at end of file
diff --git a/aiak_training_llm/models/fastvit/mobileclip/configs/mobileclip_l.json b/aiak_training_llm/models/fastvit/mobileclip/configs/mobileclip_l.json
new file mode 100644
index 00000000..c6f98c79
--- /dev/null
+++ b/aiak_training_llm/models/fastvit/mobileclip/configs/mobileclip_l.json
@@ -0,0 +1,20 @@
+{
+ "embed_dim": 768,
+ "image_cfg": {
+ "image_size": 1024,
+ "model_name": "fastvithd",
+ "embed_dim": 3072,
+ "patch_size": 64
+ },
+ "text_cfg": {
+ "context_length": 77,
+ "vocab_size": 49408,
+ "dim": 768,
+ "ffn_multiplier_per_layer": 4.0,
+ "n_heads_per_layer": 12,
+ "n_transformer_layers": 12,
+ "norm_layer": "layer_norm_fp32",
+ "causal_masking": false,
+ "model_name": "base"
+ }
+}
\ No newline at end of file
diff --git a/aiak_training_llm/models/fastvit/mobileclip/mci.py b/aiak_training_llm/models/fastvit/mobileclip/mci.py
new file mode 100644
index 00000000..a5a61d6b
--- /dev/null
+++ b/aiak_training_llm/models/fastvit/mobileclip/mci.py
@@ -0,0 +1,1478 @@
+#
+# For licensing see accompanying LICENSE file.
+# Copyright (C) 2025 Apple Inc. All Rights Reserved.
+#
+import copy
+from functools import partial
+from typing import List, Tuple, Optional, Union, Dict
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+import torch.nn.functional as F
+from torch.nn.init import normal_
+
+from timm.models import register_model
+from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
+from timm.layers import DropPath, SqueezeExcite
+
+
+def _cfg(url="", **kwargs):
+ return {
+ "url": url,
+ "num_classes": 1000,
+ "input_size": (3, 256, 256),
+ "pool_size": None,
+ "crop_pct": 0.95,
+ "interpolation": "bicubic",
+ "mean": IMAGENET_DEFAULT_MEAN,
+ "std": IMAGENET_DEFAULT_STD,
+ "classifier": "head",
+ **kwargs,
+ }
+
+
+default_cfgs = {
+ "fastvit_t": _cfg(crop_pct=0.9),
+ "fastvit_s": _cfg(crop_pct=0.9),
+ "fastvit_m": _cfg(crop_pct=0.95),
+}
+
+
+class SEBlock(nn.Module):
+ """Squeeze and Excite module.
+
+ Pytorch implementation of `Squeeze-and-Excitation Networks` -
+ https://arxiv.org/pdf/1709.01507.pdf
+ """
+
+ def __init__(self, in_channels: int, rd_ratio: float = 0.0625) -> None:
+ """Construct a Squeeze and Excite Module.
+
+ Args:
+ in_channels: Number of input channels.
+ rd_ratio: Input channel reduction ratio.
+ """
+ super(SEBlock, self).__init__()
+ self.reduce = nn.Conv2d(
+ in_channels=in_channels,
+ out_channels=int(in_channels * rd_ratio),
+ kernel_size=1,
+ stride=1,
+ bias=True,
+ )
+ self.expand = nn.Conv2d(
+ in_channels=int(in_channels * rd_ratio),
+ out_channels=in_channels,
+ kernel_size=1,
+ stride=1,
+ bias=True,
+ )
+
+ def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+ """Apply forward pass."""
+ b, c, h, w = inputs.size()
+ x = F.avg_pool2d(inputs, kernel_size=[h, w])
+ x = self.reduce(x)
+ x = F.relu(x)
+ x = self.expand(x)
+ x = torch.sigmoid(x)
+ x = x.view(-1, c, 1, 1)
+ return inputs * x
+
+
+class MobileOneBlock(nn.Module):
+ """MobileOne building block.
+
+ This block has a multi-branched architecture at train-time
+ and plain-CNN style architecture at inference time
+ For more details, please refer to our paper:
+ `An Improved One millisecond Mobile Backbone` -
+ https://arxiv.org/pdf/2206.04040.pdf
+ """
+
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int,
+ kernel_size: int,
+ stride: int = 1,
+ padding: int = 0,
+ dilation: int = 1,
+ groups: int = 1,
+ inference_mode: bool = False,
+ use_se: bool = False,
+ use_act: bool = True,
+ use_scale_branch: bool = True,
+ num_conv_branches: int = 1,
+ activation: nn.Module = nn.GELU(),
+ ) -> None:
+ """Construct a MobileOneBlock module.
+
+ Args:
+ in_channels: Number of channels in the input.
+ out_channels: Number of channels produced by the block.
+ kernel_size: Size of the convolution kernel.
+ stride: Stride size.
+ padding: Zero-padding size.
+ dilation: Kernel dilation factor.
+ groups: Group number.
+ inference_mode: If True, instantiates model in inference mode.
+ use_se: Whether to use SE-ReLU activations.
+ use_act: Whether to use activation. Default: ``True``
+ use_scale_branch: Whether to use scale branch. Default: ``True``
+ num_conv_branches: Number of linear conv branches.
+ """
+ super(MobileOneBlock, self).__init__()
+ self.inference_mode = inference_mode
+ self.groups = groups
+ self.stride = stride
+ self.padding = padding
+ self.dilation = dilation
+ self.kernel_size = kernel_size
+ self.in_channels = in_channels
+ self.out_channels = out_channels
+ self.num_conv_branches = num_conv_branches
+
+ # Check if SE-ReLU is requested
+ if use_se:
+ self.se = SEBlock(out_channels)
+ else:
+ self.se = nn.Identity()
+
+ if use_act:
+ self.activation = activation
+ else:
+ self.activation = nn.Identity()
+
+ if inference_mode:
+ self.reparam_conv = nn.Conv2d(
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=kernel_size,
+ stride=stride,
+ padding=padding,
+ dilation=dilation,
+ groups=groups,
+ bias=True,
+ )
+ else:
+ # Re-parameterizable skip connection
+ # Fallback, sometimes batchnorm tensors
+ # do not get instantiated correctly on some processes
+ # when using deepspeed + accelerate
+ norm_layer = nn.BatchNorm2d(num_features=in_channels)
+ if norm_layer.weight.shape[0] == 0:
+ norm_layer.weight = nn.Parameter(torch.zeros(in_channels))
+ if norm_layer.bias.shape[0] == 0:
+ norm_layer.bias = nn.Parameter(torch.zeros(in_channels))
+
+ self.rbr_skip = (
+ norm_layer
+ if out_channels == in_channels and stride == 1
+ else None
+ )
+
+ # Re-parameterizable conv branches
+ if num_conv_branches > 0:
+ rbr_conv = list()
+ for _ in range(self.num_conv_branches):
+ rbr_conv.append(
+ self._conv_bn(kernel_size=kernel_size, padding=padding)
+ )
+ self.rbr_conv = nn.ModuleList(rbr_conv)
+ else:
+ self.rbr_conv = None
+
+ # Re-parameterizable scale branch
+ self.rbr_scale = None
+ if not isinstance(kernel_size, int):
+ kernel_size = kernel_size[0]
+ if (kernel_size > 1) and use_scale_branch:
+ self.rbr_scale = self._conv_bn(kernel_size=1, padding=0)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ """Apply forward pass."""
+ # Inference mode forward pass.
+ if self.inference_mode:
+ return self.activation(self.se(self.reparam_conv(x)))
+
+ # Multi-branched train-time forward pass.
+ # Skip branch output
+ identity_out = 0
+ if self.rbr_skip is not None:
+ identity_out = self.rbr_skip(x)
+
+ # Scale branch output
+ scale_out = 0
+ if self.rbr_scale is not None:
+ scale_out = self.rbr_scale(x)
+
+ # Other branches
+ out = scale_out + identity_out
+ if self.rbr_conv is not None:
+ for ix in range(self.num_conv_branches):
+ out += self.rbr_conv[ix](x)
+
+ return self.activation(self.se(out))
+
+ def reparameterize(self):
+ """Following works like `RepVGG: Making VGG-style ConvNets Great Again` -
+ https://arxiv.org/pdf/2101.03697.pdf. We re-parameterize multi-branched
+ architecture used at training time to obtain a plain CNN-like structure
+ for inference.
+ """
+ if self.inference_mode:
+ return
+ kernel, bias = self._get_kernel_bias()
+ self.reparam_conv = nn.Conv2d(
+ in_channels=self.in_channels,
+ out_channels=self.out_channels,
+ kernel_size=self.kernel_size,
+ stride=self.stride,
+ padding=self.padding,
+ dilation=self.dilation,
+ groups=self.groups,
+ bias=True,
+ )
+ self.reparam_conv.weight.data = kernel
+ self.reparam_conv.bias.data = bias
+
+ # Delete un-used branches
+ self.__delattr__("rbr_conv")
+ self.__delattr__("rbr_scale")
+ if hasattr(self, "rbr_skip"):
+ self.__delattr__("rbr_skip")
+
+ self.inference_mode = True
+
+ def _get_kernel_bias(self) -> Tuple[torch.Tensor, torch.Tensor]:
+ """Method to obtain re-parameterized kernel and bias.
+ Reference: https://github.com/DingXiaoH/RepVGG/blob/main/repvgg.py#L83
+
+ Returns:
+ Tuple of (kernel, bias) after fusing branches.
+ """
+ # get weights and bias of scale branch
+ kernel_scale = 0
+ bias_scale = 0
+ if self.rbr_scale is not None:
+ kernel_scale, bias_scale = self._fuse_bn_tensor(self.rbr_scale)
+ # Pad scale branch kernel to match conv branch kernel size.
+ pad = self.kernel_size // 2
+ kernel_scale = torch.nn.functional.pad(kernel_scale, [pad, pad, pad, pad])
+
+ # get weights and bias of skip branch
+ kernel_identity = 0
+ bias_identity = 0
+ if self.rbr_skip is not None:
+ kernel_identity, bias_identity = self._fuse_bn_tensor(self.rbr_skip)
+
+ # get weights and bias of conv branches
+ kernel_conv = 0
+ bias_conv = 0
+ if self.rbr_conv is not None:
+ for ix in range(self.num_conv_branches):
+ _kernel, _bias = self._fuse_bn_tensor(self.rbr_conv[ix])
+ kernel_conv += _kernel
+ bias_conv += _bias
+
+ kernel_final = kernel_conv + kernel_scale + kernel_identity
+ bias_final = bias_conv + bias_scale + bias_identity
+ return kernel_final, bias_final
+
+ def _fuse_bn_tensor(
+ self, branch: Union[nn.Sequential, nn.BatchNorm2d]
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """Method to fuse batchnorm layer with preceeding conv layer.
+ Reference: https://github.com/DingXiaoH/RepVGG/blob/main/repvgg.py#L95
+
+ Args:
+ branch: Sequence of ops to be fused.
+
+ Returns:
+ Tuple of (kernel, bias) after fusing batchnorm.
+ """
+ if isinstance(branch, nn.Sequential):
+ kernel = branch.conv.weight
+ running_mean = branch.bn.running_mean
+ running_var = branch.bn.running_var
+ gamma = branch.bn.weight
+ beta = branch.bn.bias
+ eps = branch.bn.eps
+ else:
+ assert isinstance(branch, nn.BatchNorm2d)
+ if not hasattr(self, "id_tensor"):
+ input_dim = self.in_channels // self.groups
+
+ kernel_size = self.kernel_size
+ if isinstance(self.kernel_size, int):
+ kernel_size = (self.kernel_size, self.kernel_size)
+
+ kernel_value = torch.zeros(
+ (self.in_channels, input_dim, kernel_size[0], kernel_size[1]),
+ dtype=branch.weight.dtype,
+ device=branch.weight.device,
+ )
+ for i in range(self.in_channels):
+ kernel_value[
+ i, i % input_dim, kernel_size[0] // 2, kernel_size[1] // 2
+ ] = 1
+ self.id_tensor = kernel_value
+ kernel = self.id_tensor
+ running_mean = branch.running_mean
+ running_var = branch.running_var
+ gamma = branch.weight
+ beta = branch.bias
+ eps = branch.eps
+ std = (running_var + eps).sqrt()
+ t = (gamma / std).reshape(-1, 1, 1, 1)
+ return kernel * t, beta - running_mean * gamma / std
+
+ def _conv_bn(self, kernel_size: int, padding: int) -> nn.Sequential:
+ """Helper method to construct conv-batchnorm layers.
+
+ Args:
+ kernel_size: Size of the convolution kernel.
+ padding: Zero-padding size.
+
+ Returns:
+ Conv-BN module.
+ """
+ # Fallback, sometimes batchnorm tensors
+ # do not get instantiated correctly on some processes
+ # when using deepspeed + accelerate
+ norm_layer = nn.BatchNorm2d(num_features=self.out_channels)
+ if norm_layer.weight.shape[0] == 0:
+ norm_layer.weight = nn.Parameter(torch.zeros(self.out_channels))
+ if norm_layer.bias.shape[0] == 0:
+ norm_layer.bias = nn.Parameter(torch.zeros(self.out_channels))
+
+ mod_list = nn.Sequential()
+ mod_list.add_module(
+ "conv",
+ nn.Conv2d(
+ in_channels=self.in_channels,
+ out_channels=self.out_channels,
+ kernel_size=kernel_size,
+ stride=self.stride,
+ padding=padding,
+ groups=self.groups,
+ bias=False,
+ ),
+ )
+ mod_list.add_module("bn", norm_layer)
+ return mod_list
+
+
+class ReparamLargeKernelConv(nn.Module):
+ """Building Block of RepLKNet
+
+ This class defines overparameterized large kernel conv block
+ introduced in `RepLKNet `_
+
+ Reference: https://github.com/DingXiaoH/RepLKNet-pytorch
+ """
+
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int,
+ kernel_size: int,
+ stride: int,
+ groups: int,
+ small_kernel: int,
+ inference_mode: bool = False,
+ use_se: bool = False,
+ activation: nn.Module = nn.GELU(),
+ ) -> None:
+ """Construct a ReparamLargeKernelConv module.
+
+ Args:
+ in_channels: Number of input channels.
+ out_channels: Number of output channels.
+ kernel_size: Kernel size of the large kernel conv branch.
+ stride: Stride size. Default: 1
+ groups: Group number. Default: 1
+ small_kernel: Kernel size of small kernel conv branch.
+ inference_mode: If True, instantiates model in inference mode. Default: ``False``
+ activation: Activation module. Default: ``nn.GELU``
+ """
+ super(ReparamLargeKernelConv, self).__init__()
+
+ self.stride = stride
+ self.groups = groups
+ self.in_channels = in_channels
+ self.out_channels = out_channels
+ self.activation = activation
+
+ self.kernel_size = kernel_size
+ self.small_kernel = small_kernel
+ self.padding = kernel_size // 2
+
+ # Check if SE is requested
+ if use_se:
+ self.se = SqueezeExcite(out_channels, rd_ratio=0.25)
+ else:
+ self.se = nn.Identity()
+
+ if inference_mode:
+ self.lkb_reparam = nn.Conv2d(
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=kernel_size,
+ stride=stride,
+ padding=self.padding,
+ dilation=1,
+ groups=groups,
+ bias=True,
+ )
+ else:
+ self.lkb_origin = self._conv_bn(
+ kernel_size=kernel_size, padding=self.padding
+ )
+ if small_kernel is not None:
+ assert (
+ small_kernel <= kernel_size
+ ), "The kernel size for re-param cannot be larger than the large kernel!"
+ self.small_conv = self._conv_bn(
+ kernel_size=small_kernel, padding=small_kernel // 2
+ )
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ """Apply forward pass."""
+ if hasattr(self, "lkb_reparam"):
+ out = self.lkb_reparam(x)
+ else:
+ out = self.lkb_origin(x)
+ if hasattr(self, "small_conv"):
+ out += self.small_conv(x)
+
+ return self.activation(self.se(out))
+
+ def get_kernel_bias(self) -> Tuple[torch.Tensor, torch.Tensor]:
+ """Method to obtain re-parameterized kernel and bias.
+ Reference: https://github.com/DingXiaoH/RepLKNet-pytorch
+
+ Returns:
+ Tuple of (kernel, bias) after fusing branches.
+ """
+ eq_k, eq_b = self._fuse_bn(self.lkb_origin.conv, self.lkb_origin.bn)
+ if hasattr(self, "small_conv"):
+ small_k, small_b = self._fuse_bn(self.small_conv.conv, self.small_conv.bn)
+ eq_b += small_b
+ eq_k += nn.functional.pad(
+ small_k, [(self.kernel_size - self.small_kernel) // 2] * 4
+ )
+ return eq_k, eq_b
+
+ def reparameterize(self) -> None:
+ """
+ Following works like `RepVGG: Making VGG-style ConvNets Great Again` -
+ https://arxiv.org/pdf/2101.03697.pdf. We re-parameterize multi-branched
+ architecture used at training time to obtain a plain CNN-like structure
+ for inference.
+ """
+ eq_k, eq_b = self.get_kernel_bias()
+ self.lkb_reparam = nn.Conv2d(
+ in_channels=self.in_channels,
+ out_channels=self.out_channels,
+ kernel_size=self.kernel_size,
+ stride=self.stride,
+ padding=self.padding,
+ dilation=self.lkb_origin.conv.dilation,
+ groups=self.groups,
+ bias=True,
+ )
+
+ self.lkb_reparam.weight.data = eq_k
+ self.lkb_reparam.bias.data = eq_b
+ self.__delattr__("lkb_origin")
+ if hasattr(self, "small_conv"):
+ self.__delattr__("small_conv")
+
+ @staticmethod
+ def _fuse_bn(
+ conv: torch.Tensor, bn: nn.BatchNorm2d
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """Method to fuse batchnorm layer with conv layer.
+
+ Args:
+ conv: Convolutional kernel weights.
+ bn: Batchnorm 2d layer.
+
+ Returns:
+ Tuple of (kernel, bias) after fusing batchnorm.
+ """
+ kernel = conv.weight
+ running_mean = bn.running_mean
+ running_var = bn.running_var
+ gamma = bn.weight
+ beta = bn.bias
+ eps = bn.eps
+ std = (running_var + eps).sqrt()
+ t = (gamma / std).reshape(-1, 1, 1, 1)
+ return kernel * t, beta - running_mean * gamma / std
+
+ def _conv_bn(self, kernel_size: int, padding: int = 0) -> nn.Sequential:
+ """Helper method to construct conv-batchnorm layers.
+
+ Args:
+ kernel_size: Size of the convolution kernel.
+ padding: Zero-padding size.
+
+ Returns:
+ A nn.Sequential Conv-BN module.
+ """
+ # Fallback, sometimes batchnorm tensors
+ # do not get instantiated correctly on some processes
+ # when using deepspeed + accelerate
+ norm_layer = nn.BatchNorm2d(num_features=self.out_channels)
+ if norm_layer.weight.shape[0] == 0:
+ norm_layer.weight = nn.Parameter(torch.zeros(self.out_channels))
+ if norm_layer.bias.shape[0] == 0:
+ norm_layer.bias = nn.Parameter(torch.zeros(self.out_channels))
+
+ mod_list = nn.Sequential()
+ mod_list.add_module(
+ "conv",
+ nn.Conv2d(
+ in_channels=self.in_channels,
+ out_channels=self.out_channels,
+ kernel_size=kernel_size,
+ stride=self.stride,
+ padding=padding,
+ groups=self.groups,
+ bias=False,
+ ),
+ )
+ mod_list.add_module("bn", norm_layer)
+ return mod_list
+
+
+def convolutional_stem(
+ in_channels: int, out_channels: int, inference_mode: bool = False, use_scale_branch: bool = True,
+) -> nn.Sequential:
+ """Build convolutional stem with MobileOne blocks.
+
+ Args:
+ in_channels: Number of input channels.
+ out_channels: Number of output channels.
+ inference_mode: Flag to instantiate model in inference mode. Default: ``False``
+
+ Returns:
+ nn.Sequential object with stem elements.
+ """
+ return nn.Sequential(
+ MobileOneBlock(
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=3,
+ stride=2,
+ padding=1,
+ groups=1,
+ inference_mode=inference_mode,
+ use_se=False,
+ num_conv_branches=1,
+ use_scale_branch=use_scale_branch
+ ),
+ MobileOneBlock(
+ in_channels=out_channels,
+ out_channels=out_channels,
+ kernel_size=3,
+ stride=2,
+ padding=1,
+ groups=out_channels,
+ inference_mode=inference_mode,
+ use_se=False,
+ num_conv_branches=1,
+ use_scale_branch=use_scale_branch
+ ),
+ MobileOneBlock(
+ in_channels=out_channels,
+ out_channels=out_channels,
+ kernel_size=1,
+ stride=1,
+ padding=0,
+ groups=1,
+ inference_mode=inference_mode,
+ use_se=False,
+ num_conv_branches=1,
+ use_scale_branch=use_scale_branch
+ ),
+ )
+
+
+class LayerNormChannel(nn.Module):
+ """
+ LayerNorm only for Channel Dimension.
+ Input: tensor in shape [B, C, H, W]
+ """
+ def __init__(self, num_features, eps=1e-05) -> None:
+ super().__init__()
+ self.weight = nn.Parameter(torch.ones(num_features))
+ self.bias = nn.Parameter(torch.zeros(num_features))
+ self.eps = eps
+
+ def forward(self, x) -> torch.Tensor:
+ u = x.mean(1, keepdim=True)
+ s = (x - u).pow(2).mean(1, keepdim=True)
+ x = (x - u) / torch.sqrt(s + self.eps)
+ x = self.weight.unsqueeze(-1).unsqueeze(-1) * x \
+ + self.bias.unsqueeze(-1).unsqueeze(-1)
+ return x
+
+
+class MHSA(nn.Module):
+ """Multi-headed Self Attention module.
+
+ Source modified from:
+ https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
+ """
+
+ def __init__(
+ self,
+ dim: int,
+ head_dim: int = 32,
+ qkv_bias: bool = False,
+ attn_drop: float = 0.0,
+ proj_drop: float = 0.0,
+ ) -> None:
+ """Build MHSA module that can handle 3D or 4D input tensors.
+
+ Args:
+ dim: Number of embedding dimensions.
+ head_dim: Number of hidden dimensions per head. Default: ``32``
+ qkv_bias: Use bias or not. Default: ``False``
+ attn_drop: Dropout rate for attention tensor.
+ proj_drop: Dropout rate for projection tensor.
+ """
+ super().__init__()
+ assert dim % head_dim == 0, "dim should be divisible by head_dim"
+ self.head_dim = head_dim
+ self.num_heads = dim // head_dim
+ self.scale = head_dim**-0.5
+
+ self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+ self.attn_drop = nn.Dropout(attn_drop)
+ self.proj = nn.Linear(dim, dim)
+ self.proj_drop = nn.Dropout(proj_drop)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ shape = x.shape
+ B, C, H, W = shape
+ N = H * W
+ if len(shape) == 4:
+ x = torch.flatten(x, start_dim=2).transpose(-2, -1) # (B, N, C)
+ qkv = (
+ self.qkv(x)
+ .reshape(B, N, 3, self.num_heads, self.head_dim)
+ .permute(2, 0, 3, 1, 4)
+ )
+ q, k, v = qkv.unbind(0) # make torchscript happy (cannot use tensor as tuple)
+
+ # trick here to make q@k.t more stable
+ attn = (q * self.scale) @ k.transpose(-2, -1)
+ attn = attn.softmax(dim=-1)
+ attn = self.attn_drop(attn)
+
+ x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+ x = self.proj(x)
+ x = self.proj_drop(x)
+ if len(shape) == 4:
+ x = x.transpose(-2, -1).reshape(B, C, H, W)
+
+ return x
+
+
+class PatchEmbed(nn.Module):
+ """Convolutional patch embedding layer."""
+
+ def __init__(
+ self,
+ patch_size: int,
+ stride: int,
+ in_channels: int,
+ embed_dim: int,
+ inference_mode: bool = False,
+ use_se: bool = False,
+ ) -> None:
+ """Build patch embedding layer.
+
+ Args:
+ patch_size: Patch size for embedding computation.
+ stride: Stride for convolutional embedding layer.
+ in_channels: Number of channels of input tensor.
+ embed_dim: Number of embedding dimensions.
+ inference_mode: Flag to instantiate model in inference mode. Default: ``False``
+ use_se: If ``True`` SE block will be used.
+ """
+ super().__init__()
+ block = list()
+ block.append(
+ ReparamLargeKernelConv(
+ in_channels=in_channels,
+ out_channels=embed_dim,
+ kernel_size=patch_size,
+ stride=stride,
+ groups=in_channels,
+ small_kernel=3,
+ inference_mode=inference_mode,
+ use_se=use_se,
+ )
+ )
+ block.append(
+ MobileOneBlock(
+ in_channels=embed_dim,
+ out_channels=embed_dim,
+ kernel_size=1,
+ stride=1,
+ padding=0,
+ groups=1,
+ inference_mode=inference_mode,
+ use_se=False,
+ num_conv_branches=1,
+ )
+ )
+ self.proj = nn.Sequential(*block)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ x = self.proj(x)
+ return x
+
+
+class RepMixer(nn.Module):
+ """Reparameterizable token mixer.
+
+ For more details, please refer to our paper:
+ `FastViT: A Fast Hybrid Vision Transformer using Structural Reparameterization `_
+ """
+
+ def __init__(
+ self,
+ dim,
+ kernel_size=3,
+ use_layer_scale=True,
+ layer_scale_init_value=1e-5,
+ inference_mode: bool = False,
+ ):
+ """Build RepMixer Module.
+
+ Args:
+ dim: Input feature map dimension. :math:`C_{in}` from an expected input of size :math:`(B, C_{in}, H, W)`.
+ kernel_size: Kernel size for spatial mixing. Default: 3
+ use_layer_scale: If True, learnable layer scale is used. Default: ``True``
+ layer_scale_init_value: Initial value for layer scale. Default: 1e-5
+ inference_mode: If True, instantiates model in inference mode. Default: ``False``
+ """
+ super().__init__()
+ self.dim = dim
+ self.kernel_size = kernel_size
+ self.inference_mode = inference_mode
+
+ if inference_mode:
+ self.reparam_conv = nn.Conv2d(
+ in_channels=self.dim,
+ out_channels=self.dim,
+ kernel_size=self.kernel_size,
+ stride=1,
+ padding=self.kernel_size // 2,
+ groups=self.dim,
+ bias=True,
+ )
+ else:
+ self.norm = MobileOneBlock(
+ dim,
+ dim,
+ kernel_size,
+ padding=kernel_size // 2,
+ groups=dim,
+ use_act=False,
+ use_scale_branch=False,
+ num_conv_branches=0,
+ )
+ self.mixer = MobileOneBlock(
+ dim,
+ dim,
+ kernel_size,
+ padding=kernel_size // 2,
+ groups=dim,
+ use_act=False,
+ )
+ self.use_layer_scale = use_layer_scale
+ if use_layer_scale:
+ self.layer_scale = nn.Parameter(
+ layer_scale_init_value * torch.ones((dim, 1, 1)), requires_grad=True
+ )
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ if hasattr(self, "reparam_conv"):
+ x = self.reparam_conv(x)
+ return x
+ else:
+ if self.use_layer_scale:
+ x = x + self.layer_scale * (self.mixer(x) - self.norm(x))
+ else:
+ x = x + self.mixer(x) - self.norm(x)
+ return x
+
+ def reparameterize(self) -> None:
+ """Reparameterize mixer and norm into a single
+ convolutional layer for efficient inference.
+ """
+ if self.inference_mode:
+ return
+
+ self.mixer.reparameterize()
+ self.norm.reparameterize()
+
+ if self.use_layer_scale:
+ w = self.mixer.id_tensor + self.layer_scale.unsqueeze(-1) * (
+ self.mixer.reparam_conv.weight - self.norm.reparam_conv.weight
+ )
+ b = torch.squeeze(self.layer_scale) * (
+ self.mixer.reparam_conv.bias - self.norm.reparam_conv.bias
+ )
+ else:
+ w = (
+ self.mixer.id_tensor
+ + self.mixer.reparam_conv.weight
+ - self.norm.reparam_conv.weight
+ )
+ b = self.mixer.reparam_conv.bias - self.norm.reparam_conv.bias
+
+ self.reparam_conv = nn.Conv2d(
+ in_channels=self.dim,
+ out_channels=self.dim,
+ kernel_size=self.kernel_size,
+ stride=1,
+ padding=self.kernel_size // 2,
+ groups=self.dim,
+ bias=True,
+ )
+ self.reparam_conv.weight.data = w
+ self.reparam_conv.bias.data = b
+
+ self.__delattr__("mixer")
+ self.__delattr__("norm")
+ if self.use_layer_scale:
+ self.__delattr__("layer_scale")
+
+
+class ConvFFN(nn.Module):
+ """Convolutional FFN Module."""
+
+ def __init__(
+ self,
+ in_channels: int,
+ hidden_channels: Optional[int] = None,
+ out_channels: Optional[int] = None,
+ act_layer: nn.Module = nn.GELU,
+ drop: float = 0.0,
+ ) -> None:
+ """Build convolutional FFN module.
+
+ Args:
+ in_channels: Number of input channels.
+ hidden_channels: Number of channels after expansion. Default: None
+ out_channels: Number of output channels. Default: None
+ act_layer: Activation layer. Default: ``GELU``
+ drop: Dropout rate. Default: ``0.0``.
+ """
+ super().__init__()
+ out_channels = out_channels or in_channels
+ hidden_channels = hidden_channels or in_channels
+ self.conv = nn.Sequential()
+ self.conv.add_module(
+ "conv",
+ nn.Conv2d(
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=7,
+ padding=3,
+ groups=in_channels,
+ bias=False,
+ ),
+ )
+
+ # Fallback, sometimes batchnorm tensors
+ # do not get instantiated correctly on some processes
+ # when using deepspeed + accelerate
+ norm_layer = nn.BatchNorm2d(num_features=out_channels)
+ if norm_layer.weight.shape[0] == 0:
+ norm_layer.weight = nn.Parameter(torch.zeros(out_channels))
+ if norm_layer.bias.shape[0] == 0:
+ norm_layer.bias = nn.Parameter(torch.zeros(out_channels))
+
+ self.conv.add_module("bn", norm_layer)
+ self.fc1 = nn.Conv2d(in_channels, hidden_channels, kernel_size=1)
+ self.act = act_layer()
+ self.fc2 = nn.Conv2d(hidden_channels, out_channels, kernel_size=1)
+ self.drop = nn.Dropout(drop)
+ self.apply(self._init_weights)
+
+ def _init_weights(self, m: nn.Module) -> None:
+ if isinstance(m, nn.Conv2d):
+ normal_(m.weight, std=0.02)
+ if m.bias is not None:
+ nn.init.constant_(m.bias, 0)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ x = self.conv(x)
+ x = self.fc1(x)
+ x = self.act(x)
+ x = self.drop(x)
+ x = self.fc2(x)
+ x = self.drop(x)
+ return x
+
+
+class RepCPE(nn.Module):
+ """Implementation of conditional positional encoding.
+
+ For more details refer to paper:
+ `Conditional Positional Encodings for Vision Transformers `_
+
+ In our implementation, we can reparameterize this module to eliminate a skip connection.
+ """
+
+ def __init__(
+ self,
+ in_channels: int,
+ embed_dim: int = 768,
+ spatial_shape: Union[int, Tuple[int, int]] = (7, 7),
+ inference_mode=False,
+ ) -> None:
+ """Build reparameterizable conditional positional encoding
+
+ Args:
+ in_channels: Number of input channels.
+ embed_dim: Number of embedding dimensions. Default: 768
+ spatial_shape: Spatial shape of kernel for positional encoding. Default: (7, 7)
+ inference_mode: Flag to instantiate block in inference mode. Default: ``False``
+ """
+ super(RepCPE, self).__init__()
+ if isinstance(spatial_shape, int):
+ spatial_shape = tuple([spatial_shape] * 2)
+ assert isinstance(spatial_shape, Tuple), (
+ f'"spatial_shape" must by a sequence or int, '
+ f"get {type(spatial_shape)} instead."
+ )
+ assert len(spatial_shape) == 2, (
+ f'Length of "spatial_shape" should be 2, '
+ f"got {len(spatial_shape)} instead."
+ )
+
+ self.spatial_shape = spatial_shape
+ self.embed_dim = embed_dim
+ self.in_channels = in_channels
+ self.groups = embed_dim
+
+ if inference_mode:
+ self.reparam_conv = nn.Conv2d(
+ in_channels=self.in_channels,
+ out_channels=self.embed_dim,
+ kernel_size=self.spatial_shape,
+ stride=1,
+ padding=int(self.spatial_shape[0] // 2),
+ groups=self.embed_dim,
+ bias=True,
+ )
+ else:
+ self.pe = nn.Conv2d(
+ in_channels,
+ embed_dim,
+ spatial_shape,
+ 1,
+ int(spatial_shape[0] // 2),
+ bias=True,
+ groups=embed_dim,
+ )
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ if hasattr(self, "reparam_conv"):
+ x = self.reparam_conv(x)
+ return x
+ else:
+ x = self.pe(x) + x
+ return x
+
+ def reparameterize(self) -> None:
+ # Build equivalent Id tensor
+ input_dim = self.in_channels // self.groups
+ kernel_value = torch.zeros(
+ (
+ self.in_channels,
+ input_dim,
+ self.spatial_shape[0],
+ self.spatial_shape[1],
+ ),
+ dtype=self.pe.weight.dtype,
+ device=self.pe.weight.device,
+ )
+ for i in range(self.in_channels):
+ kernel_value[
+ i,
+ i % input_dim,
+ self.spatial_shape[0] // 2,
+ self.spatial_shape[1] // 2,
+ ] = 1
+ id_tensor = kernel_value
+
+ # Reparameterize Id tensor and conv
+ w_final = id_tensor + self.pe.weight
+ b_final = self.pe.bias
+
+ # Introduce reparam conv
+ self.reparam_conv = nn.Conv2d(
+ in_channels=self.in_channels,
+ out_channels=self.embed_dim,
+ kernel_size=self.spatial_shape,
+ stride=1,
+ padding=int(self.spatial_shape[0] // 2),
+ groups=self.embed_dim,
+ bias=True,
+ )
+ self.reparam_conv.weight.data = w_final
+ self.reparam_conv.bias.data = b_final
+
+ self.__delattr__("pe")
+
+
+class RepMixerBlock(nn.Module):
+ """Implementation of Metaformer block with RepMixer as token mixer.
+
+ For more details on Metaformer structure, please refer to:
+ `MetaFormer Is Actually What You Need for Vision `_
+ """
+
+ def __init__(
+ self,
+ dim: int,
+ kernel_size: int = 3,
+ mlp_ratio: float = 4.0,
+ act_layer: nn.Module = nn.GELU,
+ drop: float = 0.0,
+ drop_path: float = 0.0,
+ use_layer_scale: bool = True,
+ layer_scale_init_value: float = 1e-5,
+ inference_mode: bool = False,
+ ):
+ """Build RepMixer Block.
+
+ Args:
+ dim: Number of embedding dimensions.
+ kernel_size: Kernel size for repmixer. Default: 3
+ mlp_ratio: MLP expansion ratio. Default: 4.0
+ act_layer: Activation layer. Default: ``nn.GELU``
+ drop: Dropout rate. Default: 0.0
+ drop_path: Drop path rate. Default: 0.0
+ use_layer_scale: Flag to turn on layer scale. Default: ``True``
+ layer_scale_init_value: Layer scale value at initialization. Default: 1e-5
+ inference_mode: Flag to instantiate block in inference mode. Default: ``False``
+ """
+
+ super().__init__()
+
+ self.token_mixer = RepMixer(
+ dim,
+ kernel_size=kernel_size,
+ use_layer_scale=use_layer_scale,
+ layer_scale_init_value=layer_scale_init_value,
+ inference_mode=inference_mode,
+ )
+
+ assert mlp_ratio > 0, "MLP ratio should be greater than 0, found: {}".format(
+ mlp_ratio
+ )
+ mlp_hidden_dim = int(dim * mlp_ratio)
+ self.convffn = ConvFFN(
+ in_channels=dim,
+ hidden_channels=mlp_hidden_dim,
+ act_layer=act_layer,
+ drop=drop,
+ )
+
+ # Drop Path
+ self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+ # Layer Scale
+ self.use_layer_scale = use_layer_scale
+ if use_layer_scale:
+ self.layer_scale = nn.Parameter(
+ layer_scale_init_value * torch.ones((dim, 1, 1)), requires_grad=True
+ )
+
+ def forward(self, x):
+ if self.use_layer_scale:
+ x = self.token_mixer(x)
+ x = x + self.drop_path(self.layer_scale * self.convffn(x))
+ else:
+ x = self.token_mixer(x)
+ x = x + self.drop_path(self.convffn(x))
+ return x
+
+
+class AttentionBlock(nn.Module):
+ """Implementation of metaformer block with MHSA as token mixer.
+
+ For more details on Metaformer structure, please refer to:
+ `MetaFormer Is Actually What You Need for Vision `_
+ """
+
+ def __init__(
+ self,
+ dim: int,
+ mlp_ratio: float = 4.0,
+ act_layer: nn.Module = nn.GELU,
+ norm_layer: nn.Module = nn.BatchNorm2d,
+ drop: float = 0.0,
+ drop_path: float = 0.0,
+ use_layer_scale: bool = True,
+ layer_scale_init_value: float = 1e-5,
+ ):
+ """Build Attention Block.
+
+ Args:
+ dim: Number of embedding dimensions.
+ mlp_ratio: MLP expansion ratio. Default: 4.0
+ act_layer: Activation layer. Default: ``nn.GELU``
+ norm_layer: Normalization layer. Default: ``nn.BatchNorm2d``
+ drop: Dropout rate. Default: 0.0
+ drop_path: Drop path rate. Default: 0.0
+ use_layer_scale: Flag to turn on layer scale. Default: ``True``
+ layer_scale_init_value: Layer scale value at initialization. Default: 1e-5
+ """
+
+ super().__init__()
+
+ # Fallback, sometimes batchnorm tensors
+ # do not get instantiated correctly on some processes
+ # when using deepspeed + accelerate
+ norm_layer_ = norm_layer(num_features=dim)
+ if norm_layer_.weight.shape[0] == 0:
+ norm_layer_.weight = nn.Parameter(torch.zeros(dim))
+ if norm_layer_.bias.shape[0] == 0:
+ norm_layer_.bias = nn.Parameter(torch.zeros(dim))
+
+ self.norm = norm_layer_
+ self.token_mixer = MHSA(dim=dim)
+
+ assert mlp_ratio > 0, "MLP ratio should be greater than 0, found: {}".format(
+ mlp_ratio
+ )
+ mlp_hidden_dim = int(dim * mlp_ratio)
+ self.convffn = ConvFFN(
+ in_channels=dim,
+ hidden_channels=mlp_hidden_dim,
+ act_layer=act_layer,
+ drop=drop,
+ )
+
+ # Drop path
+ self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+ # Layer Scale
+ self.use_layer_scale = use_layer_scale
+ if use_layer_scale:
+ self.layer_scale_1 = nn.Parameter(
+ layer_scale_init_value * torch.ones((dim, 1, 1)), requires_grad=True
+ )
+ self.layer_scale_2 = nn.Parameter(
+ layer_scale_init_value * torch.ones((dim, 1, 1)), requires_grad=True
+ )
+
+ def forward(self, x):
+ if self.use_layer_scale:
+ x = x + self.drop_path(self.layer_scale_1 * self.token_mixer(self.norm(x)))
+ x = x + self.drop_path(self.layer_scale_2 * self.convffn(x))
+ else:
+ x = x + self.drop_path(self.token_mixer(self.norm(x)))
+ x = x + self.drop_path(self.convffn(x))
+ return x
+
+
+def basic_blocks(
+ dim: int,
+ block_index: int,
+ num_blocks: List[int],
+ token_mixer_type: str,
+ kernel_size: int = 3,
+ mlp_ratio: float = 4.0,
+ act_layer: nn.Module = nn.GELU,
+ norm_layer: nn.Module = nn.BatchNorm2d,
+ drop_rate: float = 0.0,
+ drop_path_rate: float = 0.0,
+ use_layer_scale: bool = True,
+ layer_scale_init_value: float = 1e-5,
+ inference_mode=False,
+) -> nn.Sequential:
+ """Build FastViT blocks within a stage.
+
+ Args:
+ dim: Number of embedding dimensions.
+ block_index: block index.
+ num_blocks: List containing number of blocks per stage.
+ token_mixer_type: Token mixer type.
+ kernel_size: Kernel size for repmixer.
+ mlp_ratio: MLP expansion ratio.
+ act_layer: Activation layer.
+ norm_layer: Normalization layer.
+ drop_rate: Dropout rate.
+ drop_path_rate: Drop path rate.
+ use_layer_scale: Flag to turn on layer scale regularization.
+ layer_scale_init_value: Layer scale value at initialization.
+ inference_mode: Flag to instantiate block in inference mode.
+
+ Returns:
+ nn.Sequential object of all the blocks within the stage.
+ """
+ blocks = []
+ for block_idx in range(num_blocks[block_index]):
+ block_dpr = (
+ drop_path_rate
+ * (block_idx + sum(num_blocks[:block_index]))
+ / (sum(num_blocks) - 1)
+ )
+ if token_mixer_type == "repmixer":
+ blocks.append(
+ RepMixerBlock(
+ dim,
+ kernel_size=kernel_size,
+ mlp_ratio=mlp_ratio,
+ act_layer=act_layer,
+ drop=drop_rate,
+ drop_path=block_dpr,
+ use_layer_scale=use_layer_scale,
+ layer_scale_init_value=layer_scale_init_value,
+ inference_mode=inference_mode,
+ )
+ )
+ elif token_mixer_type == "attention":
+ blocks.append(
+ AttentionBlock(
+ dim,
+ mlp_ratio=mlp_ratio,
+ act_layer=act_layer,
+ norm_layer=norm_layer,
+ drop=drop_rate,
+ drop_path=block_dpr,
+ use_layer_scale=use_layer_scale,
+ layer_scale_init_value=layer_scale_init_value,
+ )
+ )
+ else:
+ raise ValueError(
+ "Token mixer type: {} not supported".format(token_mixer_type)
+ )
+ blocks = nn.Sequential(*blocks)
+ return blocks
+
+
+class GlobalPool2D(nn.Module):
+ """This class implements global pooling with linear projection."""
+
+ def __init__(self, in_dim: int, out_dim: int, *args, **kwargs) -> None:
+ super().__init__()
+ scale = in_dim**-0.5
+ self.proj = nn.Parameter(scale * torch.randn(size=(in_dim, out_dim)))
+ self.in_dim = in_dim
+ self.out_dim = out_dim
+
+ def pool(self, x) -> Tensor:
+ if x.dim() == 4:
+ dims = [-2, -1]
+ elif x.dim() == 5:
+ dims = [-3, -2, -1]
+ x = torch.mean(x, dim=dims, keepdim=False)
+ return x
+
+ def forward(self, x: Tensor, *args, **kwargs) -> Tensor:
+ # x is of shape [batch, in_dim]
+ assert (
+ x.dim() == 4
+ ), "Input should be 4-dimensional (Batch x in_dim x in_height x in_width). Got: {}".format(
+ x.shape
+ )
+
+ # [batch, in_dim, in_height, in_width] --> [batch, in_dim]
+ x = self.pool(x)
+ # [batch, in_dim] x [in_dim, out_dim] --> [batch, out_dim]
+ x = x @ self.proj
+ return x
+
+
+class FastViT(nn.Module):
+ """
+ This class implements `FastViT architecture `_
+ """
+
+ def __init__(
+ self,
+ layers,
+ token_mixers: Tuple[str, ...],
+ embed_dims=None,
+ mlp_ratios=None,
+ downsamples=None,
+ se_downsamples=None,
+ repmixer_kernel_size=3,
+ norm_layer: nn.Module = nn.BatchNorm2d,
+ act_layer: nn.Module = nn.GELU,
+ num_classes=1000,
+ pos_embs=None,
+ down_patch_size=7,
+ down_stride=2,
+ drop_rate=0.0,
+ drop_path_rate=0.0,
+ use_layer_scale=True,
+ layer_scale_init_value=1e-5,
+ init_cfg=None,
+ pretrained=None,
+ cls_ratio=2.0,
+ inference_mode=False,
+ stem_scale_branch=True,
+ **kwargs,
+ ) -> None:
+
+ super().__init__()
+
+ self.num_classes = num_classes
+ if len(layers) == 4:
+ self.out_indices = [0, 2, 4, 7]
+ elif len(layers) == 5:
+ self.out_indices = [0, 2, 4, 7, 10]
+ else:
+ raise NotImplementedError("FPN is not implemented for more than 5 stages.")
+
+ if pos_embs is None:
+ pos_embs = [None] * len(layers)
+
+ if se_downsamples is None:
+ se_downsamples = [False] * len(layers)
+
+ # Convolutional stem
+ self.patch_embed = convolutional_stem(3, embed_dims[0], inference_mode,
+ use_scale_branch=stem_scale_branch)
+
+ # Build the main stages of the network architecture
+ network = []
+ for i in range(len(layers)):
+ # Add position embeddings if requested
+ if pos_embs[i] is not None:
+ network.append(
+ pos_embs[i](
+ embed_dims[i], embed_dims[i], inference_mode=inference_mode
+ )
+ )
+ stage = basic_blocks(
+ embed_dims[i],
+ i,
+ layers,
+ token_mixer_type=token_mixers[i],
+ kernel_size=repmixer_kernel_size,
+ mlp_ratio=mlp_ratios[i],
+ act_layer=act_layer,
+ norm_layer=norm_layer,
+ drop_rate=drop_rate,
+ drop_path_rate=drop_path_rate,
+ use_layer_scale=use_layer_scale,
+ layer_scale_init_value=layer_scale_init_value,
+ inference_mode=inference_mode,
+ )
+ network.append(stage)
+ if i >= len(layers) - 1:
+ break
+
+ # Patch merging/downsampling between stages.
+ if downsamples[i] or embed_dims[i] != embed_dims[i + 1]:
+ network.append(
+ PatchEmbed(
+ patch_size=down_patch_size,
+ stride=down_stride,
+ in_channels=embed_dims[i],
+ embed_dim=embed_dims[i + 1],
+ inference_mode=inference_mode,
+ use_se=se_downsamples[i + 1],
+ )
+ )
+ self.network = nn.ModuleList(network)
+
+ # Classifier head
+ self.conv_exp = MobileOneBlock(
+ in_channels=embed_dims[-1],
+ out_channels=int(embed_dims[-1] * cls_ratio),
+ kernel_size=3,
+ stride=1,
+ padding=1,
+ groups=embed_dims[-1],
+ inference_mode=inference_mode,
+ use_se=True,
+ num_conv_branches=1,
+ )
+ self.head = (
+ nn.Linear(int(embed_dims[-1] * cls_ratio), num_classes)
+ if num_classes > 0
+ else nn.Identity()
+ )
+ self.apply(self.cls_init_weights)
+ self.init_cfg = copy.deepcopy(init_cfg)
+
+ def cls_init_weights(self, m: nn.Module) -> None:
+ """Init. for classification"""
+ if isinstance(m, nn.Linear):
+ normal_(m.weight, std=0.02)
+ if isinstance(m, nn.Linear) and m.bias is not None:
+ nn.init.constant_(m.bias, 0)
+
+ def forward_embeddings(self, x: torch.Tensor) -> torch.Tensor:
+ x = self.patch_embed(x)
+ return x
+
+ def forward_tokens(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
+ for block in self.network:
+ x = block(x)
+ return x
+
+ def forward(self, x: torch.Tensor, *args, **kwargs) -> Union[Tensor, Dict[str, Tensor]]:
+ # input embedding
+ x = self.forward_embeddings(x)
+ # through backbone
+ x = self.forward_tokens(x)
+ # for image classification/embedding
+ x = self.conv_exp(x)
+ cls_out = self.head(x)
+
+ out_dict = dict()
+ if kwargs.get("return_image_embeddings", False):
+ out_dict.update({"logits": cls_out})
+ out_dict.update({"image_embeddings": x})
+ return out_dict
+ else:
+ return cls_out
+
+
+@register_model
+def fastvithd(pretrained=False, **kwargs):
+ """Instantiate FastViTHD model variant."""
+ layers = [2, 12, 24, 4, 2]
+ embed_dims = [96, 192, 384, 768, 1536]
+ mlp_ratios = [4, 4, 4, 4, 4]
+ downsamples = [True, True, True, True, True]
+ pos_embs = [None, None, None, partial(RepCPE, spatial_shape=(7, 7)), partial(RepCPE, spatial_shape=(7, 7))]
+ token_mixers = ("repmixer", "repmixer", "repmixer", "attention", "attention")
+ model = FastViT(
+ layers,
+ token_mixers=token_mixers,
+ embed_dims=embed_dims,
+ pos_embs=pos_embs,
+ mlp_ratios=mlp_ratios,
+ downsamples=downsamples,
+ norm_layer=LayerNormChannel,
+ stem_scale_branch=False,
+ inference_mode=True,
+ **kwargs,
+ )
+ model.default_cfg = default_cfgs["fastvit_m"]
+ if pretrained:
+ raise ValueError("Functionality not implemented.")
+ return model
\ No newline at end of file
diff --git a/aiak_training_llm/models/fastvit/mobileclip_encoder.py b/aiak_training_llm/models/fastvit/mobileclip_encoder.py
new file mode 100644
index 00000000..63d01051
--- /dev/null
+++ b/aiak_training_llm/models/fastvit/mobileclip_encoder.py
@@ -0,0 +1,115 @@
+#
+# For licensing see accompanying LICENSE file.
+# Copyright (C) 2025 Apple Inc. All Rights Reserved.
+#
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers import CLIPImageProcessor
+from aiak_training_llm.models.fastvit import mobileclip
+
+
+class MobileCLIPVisionTower(nn.Module):
+ def __init__(self, vision_tower, args, delay_load=False):
+ super().__init__()
+
+ self.is_loaded = False
+ self.vision_tower_name = vision_tower
+ self.tune_vision_tower = getattr(args, 'unfreeze_mm_vision_tower', False)
+ self.input_image_size = int(vision_tower.split("_")[-1])
+
+ # Delay load is disabled for now
+ if not delay_load:
+ self.load_model()
+ elif getattr(args, 'unfreeze_mm_vision_tower', False):
+ self.load_model()
+ else:
+ model_cfg = mobileclip.load_model_config(self.vision_tower_name)
+ self.cfg_only = model_cfg
+
+ def load_model(self, device_map=None):
+ if self.is_loaded:
+ return
+
+ # Load model config
+ model_cfg = mobileclip.load_model_config(self.vision_tower_name)
+
+ # Override default image resolution
+ model_cfg["image_cfg"]["image_size"] = self.input_image_size
+
+ self.cfg_only = model_cfg
+
+ # Build HF CLIPImageProcessor with MobileCLIP parameters
+ self.image_processor = CLIPImageProcessor(crop_size={"height": model_cfg["image_cfg"]["image_size"],
+ "width": model_cfg["image_cfg"]["image_size"]},
+ image_mean=[0.0, 0.0, 0.0],
+ image_std=[1.0, 1.0, 1.0],
+ size={"shortest_edge": model_cfg["image_cfg"]["image_size"]})
+
+ # Instantiate the image encoder
+ self.vision_tower = mobileclip.MCi(model_name=model_cfg["image_cfg"]["model_name"],
+ projection_dim=model_cfg["embed_dim"])
+
+ if not self.tune_vision_tower:
+ self.vision_tower.requires_grad_(False)
+
+ self.is_loaded = True
+
+ def feature_select(self, image_forward_outs):
+ # Features from penultimate layer (conv_exp output): (B, C, H, W)
+ image_features = image_forward_outs["image_embeddings"]
+
+ # Match FastVLM: keep the spatial FastViTHD grid as visual tokens.
+ # (B, C, H, W) -> (B, H * W, C)
+ B, C, H, W = image_features.shape
+ image_features = image_features.reshape(B, C, H * W).transpose(1, 2).contiguous()
+ return image_features
+
+ def forward(self, images):
+ if self.tune_vision_tower:
+ return self.forward_images(images)
+ else:
+ with torch.no_grad():
+ return self.forward_images(images)
+
+ def forward_images(self, images):
+ if type(images) is list:
+ image_features = []
+ for image in images:
+ image_forward_out = self.vision_tower(image.to(device=self.device, dtype=self.dtype).unsqueeze(0), return_image_embeddings=True)
+ image_feature = self.feature_select(image_forward_out).to(image.dtype)
+ image_features.append(image_feature)
+ else:
+ image_forward_outs = self.vision_tower(images.to(device=self.device, dtype=self.dtype), return_image_embeddings=True)
+ image_features = self.feature_select(image_forward_outs).to(images.dtype)
+
+ return image_features
+
+ @property
+ def dummy_feature(self):
+ return torch.zeros(1, self.hidden_size, device=self.device, dtype=self.dtype)
+
+ @property
+ def dtype(self):
+ return next(self.vision_tower.parameters()).dtype
+
+ @property
+ def device(self):
+ return next(self.vision_tower.parameters()).device
+
+ @property
+ def config(self):
+ return self.cfg_only
+
+ @property
+ def hidden_size(self):
+ return self.config["image_cfg"]["embed_dim"]
+
+ @property
+ def num_patches_per_side(self):
+ return self.config["image_cfg"]["image_size"] // self.config["image_cfg"]["patch_size"]
+
+ @property
+ def num_patches(self):
+ return (self.config["image_cfg"]["image_size"] // self.config["image_cfg"]["patch_size"]) ** 2
diff --git a/aiak_training_llm/models/llavaov_1_5/__init__.py b/aiak_training_llm/models/llavaov_1_5/__init__.py
index 521aff6d..d9563d5f 100644
--- a/aiak_training_llm/models/llavaov_1_5/__init__.py
+++ b/aiak_training_llm/models/llavaov_1_5/__init__.py
@@ -1,3 +1,9 @@
-"""qwen model"""
+"""LLaVA-OneVision-1.5 model"""
-from .llavaov_1_5_model import LlavaOnevision1_5
\ No newline at end of file
+from .llavaov_1_5_model import LlavaOnevision1_5
+from . import llavaov_1_5_config
+__all__ = [
+ "LlavaOnevision1_5",
+ "get_llava_ov_mobilellm_140m_config",
+ "get_llava_ov_mobilellm_140m_fastvit_config",
+]
\ No newline at end of file
diff --git a/aiak_training_llm/models/llavaov_1_5/adapter.py b/aiak_training_llm/models/llavaov_1_5/adapter.py
deleted file mode 100644
index 57544bbc..00000000
--- a/aiak_training_llm/models/llavaov_1_5/adapter.py
+++ /dev/null
@@ -1,72 +0,0 @@
-""" Adapters """
-
-import torch
-from dataclasses import dataclass
-from typing import Union
-from megatron.core.transformer.transformer_config import TransformerConfig
-from megatron.core.transformer.module import MegatronModule
-from megatron.core.transformer.spec_utils import ModuleSpec, build_module
-
-
-@dataclass
-class AdapterSubmodules:
- """Adapter sub-modules."""
- layernorm: Union[ModuleSpec, type] = None
- linear_fc1: Union[ModuleSpec, type] = None
- linear_fc2: Union[ModuleSpec, type] = None
-
-class Adapter(MegatronModule):
- """ Adaptor """
- def __init__(self,
- config: TransformerConfig,
- submodules: AdapterSubmodules,
- input_size: int,
- output_size: int,
- spatial_merge_size: int = 2
- ) -> None:
- super().__init__(config=config)
- self.hidden_size = input_size * (spatial_merge_size**2)
-
- self.layernorm = build_module(
- submodules.layernorm,
- config=config,
- hidden_size=input_size,
- eps=config.layernorm_epsilon,
- )
-
- self.linear_fc1 = build_module(
- submodules.linear_fc1,
- self.hidden_size,
- self.hidden_size,
- config=self.config,
- init_method=self.config.init_method,
- bias=self.config.add_bias_linear,
- skip_bias_add=False,
- parallel_mode=None,
- skip_weight_param_allocation=False,
- )
-
- self.activation_func = config.activation_func
-
- self.linear_fc2 = build_module(
- submodules.linear_fc2,
- self.hidden_size,
- output_size,
- config=self.config,
- init_method=self.config.output_layer_init_method,
- bias=self.config.add_bias_linear,
- skip_bias_add=False,
- parallel_mode=None,
- skip_weight_param_allocation=False,
- )
-
- def forward(self, x: torch.Tensor, window_index: torch.LongTensor = None) -> torch.Tensor:
- """ Forward pass."""
- x = self.layernorm(x).view(-1, self.hidden_size)
- x, _ = self.linear_fc1(x)
- x = self.activation_func(x)
- x, _ = self.linear_fc2(x)
- if window_index is not None:
- reverse_indices = torch.argsort(window_index)
- x = x[reverse_indices, :].contiguous()
- return x
\ No newline at end of file
diff --git a/aiak_training_llm/models/llavaov_1_5/llavaov_1_5_config.py b/aiak_training_llm/models/llavaov_1_5/llavaov_1_5_config.py
index b88c8ce6..00afdeca 100644
--- a/aiak_training_llm/models/llavaov_1_5/llavaov_1_5_config.py
+++ b/aiak_training_llm/models/llavaov_1_5/llavaov_1_5_config.py
@@ -157,6 +157,45 @@ def llava_one_vision_1_5_14b():
)
+@register_model_config(model_family=VisionLanguageModelFamilies.LLAVA_OV_1_5, model_arch="llava-ov-mobilellm-140m")
+def llava_ov_mobilellm_140m():
+ """
+ LLaVA-OneVision-1.5 with MobileLLM-R1-140M backbone
+
+ Integrates Facebook's MobileLLM-R1-140M (140M params) as the language model
+ replacing Qwen2.5. Architecture specs from:
+ https://github.com/facebookresearch/MobileLLM-R1
+ """
+ return LlavaOnevision1_5Config(
+ # Core architecture from MobileLLM-R1-140M
+ num_layers=15, # 15 transformer layers (vs 36 in Qwen2.5-3B)
+ hidden_size=576, # 576 hidden dimension (vs 2048 in Qwen2.5-3B)
+ ffn_hidden_size=2048, # 2048 FFN dimension
+ num_attention_heads=9, # 9 attention heads
+
+ # Grouped Query Attention (GQA) - efficiency optimization
+ group_query_attention=True, # Enable GQA
+ num_query_groups=3, # 3 KV heads for GQA (9 Q heads / 3 KV heads)
+
+ # Vocabulary - MobileLLM uses Llama3 tokenizer (128,256 tokens)
+ vocab_size_in_config_file=128256, # Llama3 vocab (vs 151936 in Qwen)
+ make_vocab_size_divisible_by=128, # Keep same for efficiency
+
+ # Embeddings - MobileLLM ties embeddings (shared input/output weights)
+ untie_embeddings_and_output_weights=False, # Tied embeddings (saves params)
+
+ # Position embeddings - RoPE with extended base for long context
+ rotary_base=8000000, # 8M RoPE base (vs 1M in Qwen, supports 32k context)
+
+ # Bias settings - MobileLLM uses no bias
+ add_qkv_bias=False, # No bias in attention (efficiency)
+ qk_layernorm=True, # MobileLLM-R1 uses QK layernorm (use_qk_norm: true in config.json)
+
+ # Normalization - standard RMSNorm
+ norm_epsilon=1e-05, # 1e-5 epsilon (vs 1e-6 in Qwen)
+ )
+
+
@dataclass
class VisionConfig:
"""configuration for vision model
@@ -191,7 +230,7 @@ def get_vision_config(model_family, model_name):
""" get vision config """
config = VisionConfig(
num_layers=24,
- hidden_size=1024,
+ hidden_size=3072,
ffn_hidden_size=4096,
num_attention_heads=16,
patch_size=14,
diff --git a/aiak_training_llm/models/llavaov_1_5/llavaov_1_5_layer_spec.py b/aiak_training_llm/models/llavaov_1_5/llavaov_1_5_layer_spec.py
index 5a4ccb4b..e9b95aaf 100644
--- a/aiak_training_llm/models/llavaov_1_5/llavaov_1_5_layer_spec.py
+++ b/aiak_training_llm/models/llavaov_1_5/llavaov_1_5_layer_spec.py
@@ -1,10 +1,9 @@
import torch
-from megatron.core.extensions.transformer_engine import (
- TEDotProductAttention, TELayerNormColumnParallelLinear, TELinear,
- TERowParallelLinear)
+from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear
+# DotProductAttention is in dot_product_attention.py in this repo.
+from megatron.core.transformer.dot_product_attention import DotProductAttention
from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add
-from megatron.core.transformer.attention import (SelfAttention,
- SelfAttentionSubmodules)
+from megatron.core.transformer.attention import SelfAttention, SelfAttentionSubmodules
from megatron.core.transformer.enums import AttnMaskType
from megatron.core.transformer.identity_op import IdentityOp
from megatron.core.transformer.mlp import MLP, MLPSubmodules
@@ -17,6 +16,9 @@
from aiak_training_llm.models.custom.common.local_norm import LocalNorm
from aiak_training_llm.models.dispatch import multiacc_modules
+from aiak_training_llm.models.mobilellm.mobilellm_layer_spec import (
+ get_mobilellm_layer_with_te_spec, # noqa: F401 - re-exported for provider use
+)
from aiak_training_llm.models.qwen_vl.qwen2_vl_layer_spec import \
get_adapeter_layer_with_spec
from aiak_training_llm.utils import is_te_min_version
@@ -25,7 +27,7 @@
def get_vision_layer_with_spec() -> ModuleSpec:
- """Use this spec for an implementation using transformer, local or multi-accel engine."""
+ """Use this spec for a standard (non-Transformer-Engine) implementation."""
return ModuleSpec(
module=TransformerLayer,
submodules=TransformerLayerSubmodules(
@@ -33,9 +35,10 @@ def get_vision_layer_with_spec() -> ModuleSpec:
module=SelfAttention,
params={"attn_mask_type": AttnMaskType.no_mask},
submodules=SelfAttentionSubmodules(
- linear_qkv=TELayerNormColumnParallelLinear,
- core_attention=TEDotProductAttention,
- linear_proj=TERowParallelLinear,
+ # use standard (non-TE) parallel linears
+ linear_qkv=ColumnParallelLinear,
+ core_attention=DotProductAttention,
+ linear_proj=RowParallelLinear,
apply_rotary_fn=apply_rotary_pos_emb_vision,
),
),
@@ -43,8 +46,8 @@ def get_vision_layer_with_spec() -> ModuleSpec:
mlp=ModuleSpec(
module=MLP,
submodules=MLPSubmodules(
- linear_fc1=TELayerNormColumnParallelLinear,
- linear_fc2=TERowParallelLinear,
+ linear_fc1=ColumnParallelLinear,
+ linear_fc2=RowParallelLinear,
),
),
mlp_bda=get_bias_dropout_add,
@@ -52,6 +55,7 @@ def get_vision_layer_with_spec() -> ModuleSpec:
)
+
def _get_mlp_module_spec(
num_experts: int=None,
moe_grouped_gemm: bool=False
@@ -138,40 +142,3 @@ def get_qwen_layer_with_te_spec(config: TransformerConfig) -> ModuleSpec:
mlp_bda=multiacc_modules.get_bias_dropout_add,
)
)
-
-# def get_qwen_layer_with_spec(qk_layernorm: bool = False) -> ModuleSpec:
-# """
-# Use this spec for an implementation using transformer, local or multi-accel engine
-# """
-# return ModuleSpec(
-# module=TransformerLayer,
-# submodules=TransformerLayerSubmodules(
-# input_layernorm=IdentityOp,
-# self_attention=ModuleSpec(
-# module=SelfAttention,
-# params={"attn_mask_type": AttnMaskType.causal},
-# submodules=SelfAttentionSubmodules(
-# linear_qkv=TELayerNormColumnParallelLinear,
-# core_attention=TEDotProductAttention,
-# linear_proj=TERowParallelLinear,
-# q_layernorm=LocalNorm if qk_layernorm else IdentityOp,
-# k_layernorm=LocalNorm if qk_layernorm else IdentityOp,
-# apply_rotary_fn=multiacc_modules.apply_rotary_pos_emb,
-# ),
-# ),
-# self_attn_bda=get_bias_dropout_add,
-# pre_mlp_layernorm=IdentityOp,
-# mlp=ModuleSpec(
-# module=MLP,
-# submodules=MLPSubmodules(
-# linear_fc1=TELayerNormColumnParallelLinear,
-# linear_fc2=TERowParallelLinear,
-# ),
-# ),
-# mlp_bda=get_bias_dropout_add,
-# sharded_state_dict_keys_map={
-# 'input_layernorm.': 'self_attention.linear_qkv.layer_norm_',
-# 'pre_mlp_layernorm.': 'mlp.linear_fc1.layer_norm_',
-# },
-# ),
-# )
\ No newline at end of file
diff --git a/aiak_training_llm/models/llavaov_1_5/llavaov_1_5_model.py b/aiak_training_llm/models/llavaov_1_5/llavaov_1_5_model.py
index e831ae12..13ee59b6 100644
--- a/aiak_training_llm/models/llavaov_1_5/llavaov_1_5_model.py
+++ b/aiak_training_llm/models/llavaov_1_5/llavaov_1_5_model.py
@@ -1,4 +1,5 @@
import logging
+import os
from collections import namedtuple
from functools import partial
from typing import List, Optional
@@ -18,9 +19,21 @@
from aiak_training_llm.models.llavaov_1_5.rice_vision_model import (
RiceViTModel, VisionModel)
+from aiak_training_llm.models.fastvit import FastViTModel
from aiak_training_llm.models.qwen import QwenModel
+from aiak_training_llm.models.mobilellm import MobileLLMModel
from aiak_training_llm.models.qwen_vl.adapter import Adapter
from aiak_training_llm.models.qwen_vl.utils import get_inputs_on_this_cp_rank
+from aiak_training_llm.utils import print_rank_0
+
+
+def _verbose_model_debug() -> bool:
+ return os.getenv("VERBOSE_MODEL_DEBUG", "0").lower() in {"1", "true", "yes", "on"}
+
+
+def _debug_rank_0(*args, **kwargs) -> None:
+ if _verbose_model_debug():
+ print_rank_0(*args, **kwargs)
def _rotate_half(x):
@@ -158,23 +171,27 @@ def __init__(
self.vision_model = None
self.adapter = None
self.language_model = None
+ self._pipeline_step = 0
# define the vision model and the projection from vision model outputs to language model inputs.
if self.add_encoder:
- # if vision_config.normalization == "RMSNorm":
- self.vision_model = RiceViTModel(
+ # Use FastViT as the vision encoder (following FastVLM repo)
+ self.vision_model = FastViTModel(
vision_config,
vision_layer_spec,
)
+ # Original Rice/SigLIP encoder (commented out):
+ # if vision_config.normalization == "RMSNorm":
+ # self.vision_model = RiceViTModel(
+ # vision_config,
+ # vision_layer_spec,
+ # )
# else:
# self.vision_model = VisionModel(
# vision_config,
# vision_layer_spec,
# )
# Map (intermediate) vision model outputs to the language model input dimension.
- # from megatron.training import print_rank_0
- # print_rank_0(f"vision_config.hidden_size: {vision_config.hidden_size}")
- # print_rank_0(f"language_config.hidden_size: {language_config.hidden_size}")
self.adapter = Adapter(
adapter_config,
adapter_layer_spec,
@@ -196,23 +213,44 @@ def __init__(
# # This attribute is needed to check if an all-reduce is required
# # on the word embeddings inside `finalize_model_grads._allreduce_word_embedding_grads`.
if self.add_decoder:
- # self.rotary_emb = Qwen2VLRotaryEmbedding(
- # dim=language_config.hidden_size // language_config.num_attention_heads,
- # theta=language_rotary_base
- # )
- self.language_model = QwenModel(
- config=language_config,
- transformer_layer_spec=language_layer_spec,
- vocab_size=language_vocab_size,
- max_sequence_length=language_max_sequence_length,
- parallel_output=parallel_output,
- position_embedding_type=language_position_embedding_type,
- rotary_percent=language_rotary_percent,
- pre_process=self.pre_process,
- post_process=self.post_process,
- rotary_base=language_rotary_base,
- share_embeddings_and_output_weights=share_embeddings_and_output_weights,
+ # Dynamically select language model based on configuration
+ # MobileLLM: 15 layers, 576 hidden size
+ # Qwen2.5: 32+ layers, 3584+ hidden size
+ is_mobilellm = (
+ language_config.num_layers <= 20 and
+ language_config.hidden_size <= 1024
)
+
+ if is_mobilellm:
+ print("[INFO] Using MobileLLM as language backbone")
+ self.language_model = MobileLLMModel(
+ config=language_config,
+ transformer_layer_spec=language_layer_spec,
+ vocab_size=language_vocab_size,
+ max_sequence_length=language_max_sequence_length,
+ parallel_output=parallel_output,
+ position_embedding_type=language_position_embedding_type,
+ rotary_percent=language_rotary_percent,
+ pre_process=self.pre_process,
+ post_process=self.post_process,
+ rotary_base=language_rotary_base,
+ share_embeddings_and_output_weights=share_embeddings_and_output_weights,
+ )
+ else:
+ print("[INFO] Using Qwen model as language backbone")
+ self.language_model = QwenModel(
+ config=language_config,
+ transformer_layer_spec=language_layer_spec,
+ vocab_size=language_vocab_size,
+ max_sequence_length=language_max_sequence_length,
+ parallel_output=parallel_output,
+ position_embedding_type=language_position_embedding_type,
+ rotary_percent=language_rotary_percent,
+ pre_process=self.pre_process,
+ post_process=self.post_process,
+ rotary_base=language_rotary_base,
+ share_embeddings_and_output_weights=share_embeddings_and_output_weights,
+ )
self.share_embeddings_and_output_weights = (
self.language_model.share_embeddings_and_output_weights
)
@@ -301,17 +339,22 @@ def forward(
output (torch.Tensor): Loss of shape [b, s] if labels are provided, otherwise logits of shape
[b, s, vocab_size].
"""
- # from megatron.training import print_rank_0
- # print_rank_0(
- # f"> forward step: input_ids shape {input_ids.shape}, "
- # f"images shape {images.shape}, "
- # f"image_grid_thw shape {image_grid_thw.shape}, "
- # # f"labels shape {labels.shape}, "
- # f"attn_mask_type {attn_mask_type}, "
- # f"position_ids shape {position_ids.shape if position_ids is not None else None}"
- # )
- # print_rank_0(input_ids)
- # print_rank_0(position_ids)
+ self._pipeline_step += 1
+ step = self._pipeline_step
+ SEP = "=" * 68
+ _debug_rank_0(f"\n{SEP}")
+ _debug_rank_0(f" PIPELINE โ STEP {step} (Stage 1 Alignment: adapter only)")
+ _debug_rank_0(f"{SEP}\n")
+
+ # โโ [1/6] BATCH โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+ _debug_rank_0(f"[1/6] BATCH")
+ _debug_rank_0(f" input_ids : {tuple(input_ids.shape) if input_ids is not None else None} {input_ids.dtype if input_ids is not None else ''}")
+ _debug_rank_0(f" images : {tuple(images.shape) if images is not None else None} {images.dtype if images is not None else ''}")
+ if labels is not None and _verbose_model_debug():
+ loss_positions = (labels != -100).sum().item()
+ total_positions = labels.numel()
+ _debug_rank_0(f" labels : {tuple(labels.shape)} {labels.dtype}")
+ _debug_rank_0(f" loss_mask : {loss_positions} / {total_positions} tokens ({100*loss_positions/max(total_positions,1):.1f}%) contribute to loss")
use_inference_kv_cache = (
inference_params is not None
and "image_tokens_count" in inference_params.key_value_memory_dict
@@ -322,15 +365,41 @@ def forward(
image_embeddings = None
elif self.add_encoder:
if images is not None:
- image_embeddings, window_index = self.vision_model(images, \
- grid_thw=image_grid_thw) # [img_len, h_vision]
+ # โโ [2/6] VISION ENCODER โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+ _debug_rank_0(f"\n[2/6] VISION ENCODER [FastViT MobileCLIP-L | FROZEN]")
+ _debug_rank_0(f" in : {tuple(images.shape)} {images.dtype}")
+ image_embeddings, window_index = self.vision_model(images, grid_thw=image_grid_thw)
+ _debug_rank_0(f" out : {tuple(image_embeddings.shape)} {image_embeddings.dtype} grad={image_embeddings.requires_grad}")
+
+ # โโ [3/6] ADAPTER โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+ _debug_rank_0(f"\n[3/6] ADAPTER [2-layer MLP {image_embeddings.shape[-1]}โ? | TRAINABLE]")
+ _debug_rank_0(f" in : {tuple(image_embeddings.shape)}")
image_embeddings = self.adapter(image_embeddings, window_index)
+ _debug_rank_0(f" out : {tuple(image_embeddings.shape)} grad={image_embeddings.requires_grad}")
+
n_image_tokens = (input_ids == self.config.image_token_id).sum().item()
n_image_features = image_embeddings.shape[0]
if n_image_tokens != n_image_features:
- raise ValueError(
- f"Image features {n_image_features} != image tokens {n_image_tokens}"
- )
+ # raise ValueError(
+ # f"Image features {n_image_features} != image tokens {n_image_tokens}"
+ # )
+ if n_image_features > n_image_tokens:
+ # Trim extra vision features if the model produced more than requested tokens.
+ image_embeddings = image_embeddings[:n_image_tokens]
+ if window_index is not None:
+ window_index = window_index[:n_image_tokens]
+ else:
+ # Pad missing vision features with zeros to align with token count.
+ pad_len = n_image_tokens - n_image_features
+ pad_emb = torch.zeros(
+ (pad_len, image_embeddings.size(1)),
+ device=image_embeddings.device,
+ dtype=image_embeddings.dtype,
+ )
+ image_embeddings = torch.cat([image_embeddings, pad_emb], dim=0)
+ if window_index is not None:
+ pad_idx = window_index[-1:].repeat(pad_len, *([1] * (window_index.dim() - 1)))
+ window_index = torch.cat([window_index, pad_idx], dim=0)
# If running inference, the language model KV cache will be updated for image token positions.
# Here we store the image tokens sequence length, which can be used as an offset to the KV cache later.
@@ -374,16 +443,20 @@ def forward(
return vision_embeddings
if self.pre_process:
+ # โโ [4/6] TOKEN FUSION โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+ _debug_rank_0(f"\n[4/6] TOKEN FUSION")
language_embeddings = self.language_model.embedding(
input_ids=input_ids, position_ids=None
) # [text_seq_len, b, h_language]
+ _debug_rank_0(f" text embeddings : {tuple(language_embeddings.shape)} {language_embeddings.dtype} grad={language_embeddings.requires_grad}")
# If running inference, we can skip image token computation if they were computed already
# earlier for this sample.
if use_inference_kv_cache or (images is None and pixel_values_videos is None):
combined_embeddings = language_embeddings
else:
- if images is not None and self.config.image_token_id in input_ids:
+ combined_embeddings = language_embeddings
+ if images is not None and (input_ids == self.config.image_token_id).any():
image_token_id = self.config.image_token_id
images_mask = (
(input_ids == image_token_id).transpose(0, 1)
@@ -392,9 +465,12 @@ def forward(
.to(language_embeddings.device)
)
image_embeddings = image_embeddings.to(language_embeddings.device, language_embeddings.dtype)
- combined_embeddings = language_embeddings.masked_scatter(images_mask, image_embeddings)
+ combined_embeddings = combined_embeddings.masked_scatter(images_mask, image_embeddings)
+ n_slots_filled = int(images_mask[..., 0].sum().item())
+ _debug_rank_0(f" image slots : {n_slots_filled} tokens replaced with vision embeddings")
+ _debug_rank_0(f" combined : {tuple(combined_embeddings.shape)} grad={combined_embeddings.requires_grad}")
- if pixel_values_videos is not None and self.config.video_token_id in input_ids:
+ if pixel_values_videos is not None and (input_ids == self.config.video_token_id).any():
video_token_id = self.config.video_token_id
videos_mask = (
(input_ids == video_token_id).transpose(0, 1)
@@ -403,7 +479,7 @@ def forward(
.to(language_embeddings.device)
)
video_embeddings = video_embeddings.to(language_embeddings.device, language_embeddings.dtype)
- combined_embeddings = language_embeddings.masked_scatter(videos_mask, video_embeddings)
+ combined_embeddings = combined_embeddings.masked_scatter(videos_mask, video_embeddings)
if self.config.context_parallel_size > 1:
combined_embeddings = get_inputs_on_this_cp_rank(combined_embeddings)
@@ -414,6 +490,10 @@ def forward(
# rotary_pos_emb = self.rotary_emb(position_ids).transpose(0, 2).contiguous()
+ # โโ [5/6] LANGUAGE MODEL โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+ _debug_rank_0(f"\n[5/6] LANGUAGE MODEL [MobileLLM-R1-140M | FROZEN]")
+ if combined_embeddings is not None:
+ _debug_rank_0(f" in : {tuple(combined_embeddings.shape)} {combined_embeddings.dtype}")
output = self.language_model(
input_ids=None,
position_ids=None,
@@ -427,7 +507,60 @@ def forward(
packed_seq_params=packed_seq_params,
extra_block_kwargs={},
)
+ out_shape = tuple(output.shape) if hasattr(output, 'shape') else None
+ _debug_rank_0(f" out : {out_shape} {getattr(output, 'dtype', None)} grad={getattr(output, 'requires_grad', None)}")
+
+ # โโ [6/6] LOSS / LOGITS โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+ _debug_rank_0(f"\n[6/6] LOSS / LOGITS")
+ if labels is not None and _verbose_model_debug():
+ # output is per-token loss with the same layout as labels [b, s]
+ loss_mask = (labels != -100)
+ if loss_mask.any():
+ # Ensure same shape as labels for boolean indexing
+ loss_tensor = output
+ if loss_tensor.shape != labels.shape:
+ # Megatron may return [s, b]; transpose to [b, s]
+ loss_tensor = loss_tensor.transpose(0, 1).contiguous()
+ valid_loss = loss_tensor[loss_mask]
+ mean_loss = valid_loss.mean().item()
+ _debug_rank_0(f" loss (mean over {loss_mask.sum().item()} tokens) : {mean_loss:.4f}")
+
+ # Optional debug-only top-1 token accuracy. This performs an
+ # extra language-model forward pass, so keep it out of normal
+ # training logs.
+ if _verbose_model_debug():
+ try:
+ with torch.no_grad():
+ _emb = combined_embeddings.detach() if combined_embeddings is not None else None
+ logits = self.language_model(
+ input_ids=None,
+ position_ids=None,
+ attention_mask=attention_mask,
+ attn_mask_type=attn_mask_type,
+ decoder_input=_emb,
+ labels=None,
+ rotary_pos_emb=None,
+ inference_params=None,
+ packed_seq_params=packed_seq_params,
+ extra_block_kwargs={},
+ ) # [s, b, vocab] or [b, s, vocab]
+ # Normalise to [b, s, vocab]
+ if logits.dim() == 3 and logits.shape[0] != labels.shape[0]:
+ logits = logits.transpose(0, 1).contiguous() # [s,b,v] -> [b,s,v]
+ # Shift: predict token i+1 from position i
+ pred = logits[:, :-1, :].argmax(dim=-1) # [b, s-1]
+ tgt = labels[:, 1:] # [b, s-1]
+ mask = (tgt != -100)
+ if mask.any():
+ correct = ((pred == tgt) & mask).sum().item()
+ total = mask.sum().item()
+ _debug_rank_0(f" top-1 accuracy : {correct}/{total} = {100*correct/total:.1f}%")
+ except Exception as acc_err:
+ _debug_rank_0(f" top-1 accuracy : (skipped - {acc_err})")
+ else:
+ _debug_rank_0(f" returning logits : {out_shape}")
+ _debug_rank_0("")
return output
diff --git a/aiak_training_llm/models/llavaov_1_5/llavaov_1_5_provider.py b/aiak_training_llm/models/llavaov_1_5/llavaov_1_5_provider.py
index f1d83929..db4a900e 100644
--- a/aiak_training_llm/models/llavaov_1_5/llavaov_1_5_provider.py
+++ b/aiak_training_llm/models/llavaov_1_5/llavaov_1_5_provider.py
@@ -1,16 +1,18 @@
-"""qwen model provider"""
+"""LLaVA-OneVision-1.5 model provider - supports Qwen and MobileLLM backbones"""
+import os
from copy import deepcopy
from dataclasses import asdict
from aiak_training_llm.models.factory import register_model_provider
from aiak_training_llm.models.llavaov_1_5.llavaov_1_5_layer_spec import (
get_adapeter_layer_with_spec, get_qwen_layer_with_te_spec,
- get_vision_layer_with_spec)
+ get_mobilellm_layer_with_te_spec, get_vision_layer_with_spec)
from aiak_training_llm.models.llavaov_1_5.llavaov_1_5_config import (
get_adapeter_config, get_vision_config)
+# MobileLLM config is now in llavaov_1_5_config.py and applied via args
from aiak_training_llm.utils import (build_transformer_config, get_args,
- print_rank_0)
+ print_rank_0, get_tokenizer)
from aiak_training_llm.utils.constants import VisionLanguageModelFamilies
from megatron.core import mpu
from megatron.core.transformer.spec_utils import import_module
@@ -18,6 +20,16 @@
from .llavaov_1_5_model import LlavaOnevision1_5
+def _verbose_model_debug() -> bool:
+ return os.getenv("VERBOSE_MODEL_DEBUG", "0").lower() in {"1", "true", "yes", "on"}
+
+
+def _debug_rank_0(*args, **kwargs) -> None:
+ if _verbose_model_debug():
+ print_rank_0(*args, **kwargs)
+
+
+# model provider registration
@register_model_provider(model_family=[VisionLanguageModelFamilies.LLAVA_OV_1_5])
def rice_vl_model_provider(
pre_process: bool = True,
@@ -40,33 +52,155 @@ def rice_vl_model_provider(
args = get_args()
print_rank_0(f'building {args.model_name} model ...')
+ _debug_rank_0(f'[DEBUG PROVIDER] Model name: {args.model_name}')
+ _debug_rank_0(f'[DEBUG PROVIDER] Args num_layers: {args.num_layers}')
+ _debug_rank_0(f'[DEBUG PROVIDER] Args hidden_size: {args.hidden_size}')
+ _debug_rank_0(f'[DEBUG PROVIDER] Args vocab_size: {args.vocab_size_in_config_file}')
+
+ config = build_transformer_config(args) #base transformer config with all hyperparams
+ _debug_rank_0(f'[DEBUG PROVIDER] TransformerConfig built with num_layers: {config.num_layers}, hidden_size: {config.hidden_size}')
- config = build_transformer_config(args)
+ language_config = deepcopy(config) # For Qwen2.5 language model
+ vision_config = deepcopy(config) # For vision encoder (SigLIP)
+ adapter_config = deepcopy(config) ## For adapter (projection)
+ _debug_rank_0(f'[DEBUG PROVIDER] Initial language_config: layers={language_config.num_layers}, hidden={language_config.hidden_size}')
- language_config = deepcopy(config)
- vision_config = deepcopy(config)
- adapter_config = deepcopy(config)
+ # Vision Encoder โ Adapter โ Language Model
+ # (SigLIP) (Projection) (Qwen2.5)
from aiak_training_llm.models import get_model_family
model_family = get_model_family(args.model_name)
- for k, v in asdict(get_vision_config(model_family, args.model_name)).items():
- setattr(vision_config, k, v)
+ _debug_rank_0(f'[DEBUG PROVIDER] Model family: {model_family}')
+
+ # Detect if using MobileLLM backbone
+ use_mobilellm = "mobilellm" in args.model_name.lower()
+ _debug_rank_0(f'[DEBUG PROVIDER] use_mobilellm flag: {use_mobilellm}')
+
+ if use_mobilellm:
+ _debug_rank_0(f'[DEBUG PROVIDER] โ Using MobileLLM-R1-140M as language backbone')
+ _debug_rank_0(f'[DEBUG PROVIDER] Language config BEFORE override: layers={language_config.num_layers}, '
+ f'hidden={language_config.hidden_size}, heads={language_config.num_attention_heads}')
+ # MobileLLM params are already in language_config from args!
+ # No need to load separately - they were applied in _validate_extra_model_args
+ _debug_rank_0(f'[DEBUG PROVIDER] Language config AFTER (should be same): layers={language_config.num_layers}, '
+ f'hidden={language_config.hidden_size}, heads={language_config.num_attention_heads}, '
+ f'query_groups={language_config.num_query_groups}')
+ else:
+ _debug_rank_0(f'[DEBUG PROVIDER] Using Qwen2.5 as language backbone')
+
+ _debug_rank_0(f'[DEBUG PROVIDER] ========== LANGUAGE CONFIG ==========')
+ _debug_rank_0(f'{language_config}')
+ _debug_rank_0(f'[DEBUG PROVIDER] ======================================')
+
+ # get vision specific config : no. of layers, hidden size, Patch size, Image resolution
+ _debug_rank_0(f'[DEBUG PROVIDER] Loading vision config...')
+
+ # Check if using FastViT - it has its own config system
+ if getattr(args, 'use_fastvit', False):
+ # FastViT loads config from mobileclip_l.json - set TransformerConfig to match
+ import json
+
+ vision_tower_name = getattr(args, 'vision_tower_name', 'mobileclip_l_1024')
+ setattr(vision_config, 'vision_tower_name', vision_tower_name)
+
+ # Parse resolution from vision_tower_name (e.g., "mobileclip_l_1024" -> 1024)
+ resolution = int(vision_tower_name.split('_')[-1]) if '_' in vision_tower_name else 1024
+
+ # Load the actual mobileclip JSON config
+ json_config_path = os.path.join(
+ os.path.dirname(__file__),
+ '../fastvit/mobileclip/configs/mobileclip_l.json'
+ )
+ with open(json_config_path, 'r') as f:
+ mobileclip_config = json.load(f)
+
+ # Extract image_cfg from the JSON
+ image_cfg = mobileclip_config['image_cfg']
+
+ # Set vision_config values from mobileclip_l.json
+ setattr(vision_config, 'num_layers', 24) # RepMixer layers in main stage
+ setattr(vision_config, 'hidden_size', image_cfg['embed_dim']) # 3072
+ setattr(vision_config, 'patch_size', image_cfg['patch_size']) # 64
+ setattr(vision_config, 'image_size', resolution) # from vision_tower_name (1024)
+ setattr(vision_config, 'num_attention_heads', image_cfg['embed_dim'] // 64) # 48
+ train_vision_model = (
+ args.trainable_modules == ['all']
+ or "vision_model" in args.trainable_modules
+ )
+ setattr(vision_config, 'unfreeze_mm_vision_tower', train_vision_model)
+
+ _debug_rank_0(f'[DEBUG PROVIDER] โ Using FastViT with vision_tower_name: {vision_tower_name}')
+ _debug_rank_0(f'[DEBUG PROVIDER] FastViT train vision tower: {train_vision_model}')
+ _debug_rank_0(f'[DEBUG PROVIDER] FastViT config loaded from {json_config_path}')
+ _debug_rank_0(f'[DEBUG PROVIDER] image_cfg: embed_dim={image_cfg["embed_dim"]}, patch_size={image_cfg["patch_size"]}, model={image_cfg["model_name"]}')
+ _debug_rank_0(f'[DEBUG PROVIDER] ========== VISION CONFIG (FastViT) ==========')
+ _debug_rank_0(f'{vision_config}')
+ _debug_rank_0(f'[DEBUG PROVIDER] ================================================')
+ else:
+ # For SigLIP/Rice models, use generic vision config
+ for k, v in asdict(get_vision_config(model_family, args.model_name)).items():
+ setattr(vision_config, k, v)
+ _debug_rank_0(f'[DEBUG PROVIDER] Vision config (SigLIP/Rice): layers={vision_config.num_layers}, hidden={vision_config.hidden_size}')
+ _debug_rank_0(f'[DEBUG PROVIDER] ========== VISION CONFIG (SigLIP/Rice) ==========')
+ _debug_rank_0(f'{vision_config}')
+ _debug_rank_0(f'[DEBUG PROVIDER] ====================================================')
+
+ # get adapter specific config : Projection dimension, Activation function
+ _debug_rank_0(f'[DEBUG PROVIDER] Loading adapter config...')
for k, v in asdict(get_adapeter_config(model_family)).items():
setattr(adapter_config, k, v)
- # print(vision_config)
- setattr(language_config, "image_token_id", 151655)
- setattr(language_config, "video_token_id", 151656)
+ _debug_rank_0(f'[DEBUG PROVIDER] ========== ADAPTER CONFIG ==========')
+ _debug_rank_0(f'{adapter_config}')
+ _debug_rank_0(f'[DEBUG PROVIDER] ======================================')
+
+ # set special token ids for language model using the shared runtime tokenizer
+ image_token_id = 151655
+ video_token_id = 151656
+ try:
+ tokenizer = get_tokenizer()
+ image_token_id = tokenizer.convert_tokens_to_ids("<|image_pad|>")
+ video_token_id = tokenizer.convert_tokens_to_ids("<|video_pad|>")
+
+ if image_token_id is None or video_token_id is None:
+ vocab = getattr(tokenizer, "vocab", None)
+ if isinstance(vocab, dict):
+ if image_token_id is None:
+ image_token_id = vocab.get("<|image_pad|>")
+ if video_token_id is None:
+ video_token_id = vocab.get("<|video_pad|>")
+
+ if image_token_id is None:
+ image_token_id = 151655
+ if video_token_id is None:
+ video_token_id = 151656
+
+ _debug_rank_0(
+ f"[DEBUG PROVIDER] Resolved vision token ids from tokenizer: "
+ f"image_token_id={image_token_id}, video_token_id={video_token_id}"
+ )
+ except Exception as e:
+ print_rank_0(
+ f"[WARN PROVIDER] Failed to resolve vision token ids from tokenizer "
+ f"({e}); fallback to defaults image_token_id={image_token_id}, video_token_id={video_token_id}"
+ )
+
+ setattr(language_config, "image_token_id", int(image_token_id))
+ setattr(language_config, "video_token_id", int(video_token_id))
+
+ #Handle pipeline parallelism
# FIXME: fix this if model_type is encoder_and_decoder
if args.encoder_pipeline_model_parallel_size in [0, None]:
+ #UNIFIED MODEL (TP=1, PP=1)
vision_config.pipeline_model_parallel_size = 1
vision_config.tensor_model_parallel_size = 1
vision_config.sequence_parallel = False
vision_config.tp_comm_overlap = False
vision_config.context_parallel_size = 1
vision_config.context_parallel_ulysses_degree = 1
- add_encoder = mpu.is_pipeline_first_stage()
- add_decoder = True
+
+ add_encoder = mpu.is_pipeline_first_stage() #True on first stage
+ add_decoder = True #Always add language model decoder
else:
assert (
args.encoder_pipeline_model_parallel_size == 1
@@ -88,11 +222,44 @@ def rice_vl_model_provider(
if args.spec is not None:
language_layer_spec = import_module(args.spec)
+ _debug_rank_0(f'[DEBUG PROVIDER] Using custom spec: {args.spec}')
else:
+ _debug_rank_0(f'[DEBUG PROVIDER] Building layer specs...')
adapter_layer_spec = get_adapeter_layer_with_spec()
vision_layer_spec = get_vision_layer_with_spec()
- language_layer_spec = get_qwen_layer_with_te_spec(language_config)
+
+ # Choose language layer spec based on backbone
+ if use_mobilellm:
+ _debug_rank_0(f'[DEBUG PROVIDER] โ Using MobileLLM layer specification')
+ _debug_rank_0(f'[DEBUG PROVIDER] MobileLLM layer config: layers={language_config.num_layers}, '
+ f'hidden={language_config.hidden_size}, heads={language_config.num_attention_heads}')
+ language_layer_spec = get_mobilellm_layer_with_te_spec(language_config)
+ _debug_rank_0(f'[DEBUG PROVIDER] MobileLLM layer spec created successfully')
+ else:
+ _debug_rank_0(f'[DEBUG PROVIDER] Using Qwen layer specification')
+ language_layer_spec = get_qwen_layer_with_te_spec(language_config)
+
+# # Vision layer spec (Transformer block)
+# - MultiheadAttention
+# - LayerNorm
+# - MLP (feedforward)
+# - Residual connections
+
+# # Language layer spec (Qwen2.5 block)
+# - MultiQueryAttention or GroupedQueryAttention
+# - RMSNorm
+# - SwiGLU MLP
+# - RoPE (Rotary Position Embedding)
+
+# # Adapter spec (projection)
+# - Linear layer(s)
+# - Optional activation
+#create the model
+ _debug_rank_0(f'[DEBUG PROVIDER] Creating LlavaOnevision1_5 model...')
+ _debug_rank_0(f'[DEBUG PROVIDER] Final language_config: layers={language_config.num_layers}, '
+ f'hidden={language_config.hidden_size}, vocab={args.padded_vocab_size}, '
+ f'rotary_base={args.rotary_base}')
model = LlavaOnevision1_5(
language_config=language_config,
vision_config=vision_config,
@@ -102,19 +269,27 @@ def rice_vl_model_provider(
adapter_layer_spec=adapter_layer_spec,
language_vocab_size=args.padded_vocab_size,
language_max_sequence_length=args.max_position_embeddings,
- pre_process=pre_process,
- post_process=post_process,
- add_encoder=add_encoder,
- add_decoder=add_decoder,
+ pre_process=pre_process, #compute embeddings?
+ post_process=post_process, #compute output logits/loss?
+ add_encoder=add_encoder, #add vision encoder?
+ add_decoder=add_decoder, #add language model decoder?
fp16_lm_cross_entropy=args.fp16_lm_cross_entropy,
parallel_output=parallel_output,
share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights,
- language_position_embedding_type=args.position_embedding_type,
+ language_position_embedding_type=args.position_embedding_type, #"rope"
language_rotary_percent=args.rotary_percent,
language_rotary_base=args.rotary_base,
seq_len_interpolation_factor=args.rotary_seq_len_interpolation_factor,
+ # When using FastViT, adapter dimensions change, so allow missing adapter weights
+ allow_missing_adapter_checkpoint=getattr(args, 'use_fastvit', False),
)
+ _debug_rank_0(f'[DEBUG PROVIDER] โ LlavaOnevision1_5 model created successfully!')
+ # Vision encoder: SigLIP with 27 layers
+ # Adapter: Projection network
+ # Language model: Qwen2.5 with 32 layers
+ # All distributed training wrappers (TP, PP, DP)
+ #freeze components if needed
if args.trainable_modules != ['all']:
train_language_model = "language_model" in args.trainable_modules
train_vision_model = "vision_model" in args.trainable_modules
@@ -122,5 +297,10 @@ def rice_vl_model_provider(
model.freeze(freeze_language_model=not train_language_model,
freeze_vision_model=not train_vision_model,
freeze_adapter=not train_adapter)
+ # Stage 0 (Pre-training): Only train adapter
+ # trainable_modules = ["adapter"]
+ # โ Freeze vision encoder โ
+ # โ Freeze language model โ
+ # โ Train adapter โ
return model
diff --git a/aiak_training_llm/models/llavaov_1_5/rice_vision_model.py b/aiak_training_llm/models/llavaov_1_5/rice_vision_model.py
index 5246b2d1..ea9ab0c9 100644
--- a/aiak_training_llm/models/llavaov_1_5/rice_vision_model.py
+++ b/aiak_training_llm/models/llavaov_1_5/rice_vision_model.py
@@ -18,19 +18,130 @@ def _rotate_half(x):
return torch.cat((-x2, x1), dim=-1)
+# def apply_rotary_pos_emb_vision(t, freqs, config, cu_seqlens=None, rotary_interleaved=False):
+# """" Apply rotation to positional embedding """
+# orig_dtype = t.dtype
+# t = t.float()
+# if cu_seqlens is not None:
+# freqs = freqs.squeeze(1)
+# cos_ = freqs.cos().float().repeat(1, 1, 2)
+# sin_ = freqs.sin().float().repeat(1, 1, 2)
+# else:
+# cos_ = freqs.cos().float().repeat(1, 1, 1, 2)
+# sin_ = freqs.sin().float().repeat(1, 1, 1, 2)
+# t = (t * cos_) + (_rotate_half(t) * sin_)
+# return t.to(orig_dtype)
+
def apply_rotary_pos_emb_vision(t, freqs, config, cu_seqlens=None, rotary_interleaved=False):
- """" Apply rotation to positional embedding """
+ """
+ Apply rotary positional embeddings to t in a robust way:
+ - build cos_/sin_ from freqs as before
+ - slice cos_/sin_ to match t sequence length
+ - expand/unsqueeze so broadcasting works for common layouts
+ """
+ # DEBUG: Check input
+ if t is None:
+ print(f"[apply_rotary_pos_emb_vision] ERROR: input t is None!", flush=True)
+ return None
+
orig_dtype = t.dtype
t = t.float()
+
+ # Build cos_/sin_ as in your original code
if cu_seqlens is not None:
- freqs = freqs.squeeze(1)
- cos_ = freqs.cos().float().repeat(1, 1, 2)
- sin_ = freqs.sin().float().repeat(1, 1, 2)
+ # freqs shape expected something like [B, 1, S, C] or similar; you squeezed axis 1
+ freqs_work = freqs.squeeze(1)
+ cos_ = freqs_work.cos().float().repeat(1, 1, 2) # -> [?, ?, 2*?] depending on freqs_work layout
+ sin_ = freqs_work.sin().float().repeat(1, 1, 2)
else:
cos_ = freqs.cos().float().repeat(1, 1, 1, 2)
sin_ = freqs.sin().float().repeat(1, 1, 1, 2)
+
+ # --- Robust alignment code: make cos_/sin_ match t's sequence axis & length ----
+ # Determine candidate sequence axis in t. Prefer axis 0, then 1, then others.
+ # We'll detect the likely seq axis by comparing sizes; fallback to axis 0.
+ def _find_seq_axis_and_len(t, cos_):
+ t_shape = tuple(t.shape)
+ # Try to find an axis i where cos_.shape[0] matches t.shape[i]
+ for axis in range(len(t_shape)):
+ if cos_.shape[0] == t_shape[axis]:
+ return axis, t_shape[axis]
+ # If no exact match, prefer axis 0 if it's smaller or cos_ is larger and can be sliced
+ return 0, t_shape[0]
+
+ seq_axis, seq_len = _find_seq_axis_and_len(t, cos_)
+
+ # Slice cos_/sin_ on their first dimension to match seq_len (safe even if cos_ is smaller)
+ if cos_.shape[0] != seq_len:
+ # if cos_ is shorter than seq_len, we'll let broadcasting fail later (better to raise)
+ if cos_.shape[0] < seq_len:
+ # Recompute or raise โ here we raise a helpful error so you can detect upstream bug
+ raise RuntimeError(
+ f"freqs/cos_ length ({cos_.shape[0]}) is smaller than the tensor sequence length ({seq_len}). "
+ "Recompute freqs for the required length or provide larger cached freqs."
+ )
+ # Slice axis 0
+ sl = [slice(None)] * cos_.ndim
+ sl[0] = slice(0, seq_len)
+ cos_ = cos_[tuple(sl)].contiguous()
+ sin_ = sin_[tuple(sl)].contiguous()
+
+ # Now make cos_/sin_ broadcast-compatible with t.
+ # Typical cases:
+ # - t: [S, B, D] and cos_: [S, D] -> unsqueeze(1) -> [S,1,D]
+ # - t: [B, S, D] and cos_: [S, D] -> unsqueeze(0) -> [1,S,D]
+ # - t: [B, H, S, D] and cos_: [S, D] -> unsqueeze to [1,1,S,1,D...] not common
+ # We'll handle the common 2/3-dim cases and otherwise attempt to add singleton dims to the left.
+ if cos_.ndim == 2 and t.ndim == 3:
+ # cos_: [S, D]
+ if seq_axis == 0:
+ # t: [S, B, D] -> need [S, 1, D]
+ cos_ = cos_.unsqueeze(1)
+ sin_ = sin_.unsqueeze(1)
+ else:
+ # seq_axis == 1: t: [B, S, D] -> need [1, S, D]
+ cos_ = cos_.unsqueeze(0)
+ sin_ = sin_.unsqueeze(0)
+ elif cos_.ndim == 2 and t.ndim == 2:
+ # both [S, D] vs [S, D] -> ok
+ pass
+ elif cos_.ndim == 3 and t.ndim == 3:
+ # cos_ might already be [S, 1, D] or [B, S, D] style; try to match by permuting if necessary
+ # If cos_.shape[0] == seq_len then assume first dim is seq and cos_ probably is [S,1,D]
+ # else attempt no-op; broadcasting should handle remaining dims
+ pass
+ else:
+ # Generic approach: try to prefix cos_ with singleton dims so that cos_.ndim == t.ndim or cos_.ndim == t.ndim-1
+ # We want sequence dim at position seq_axis in t; ensure cos_ has seq at dim 0 currently (we sliced it)
+ # So add singleton dims to the left until cos_.ndim matches t.ndim with seq at that position.
+ if cos_.ndim < t.ndim:
+ # number of dims to add on left
+ to_add = t.ndim - cos_.ndim
+ for _ in range(to_add):
+ cos_ = cos_.unsqueeze(0)
+ sin_ = sin_.unsqueeze(0)
+ # If still not aligned, rely on broadcasting but print a debug warning:
+ # (you can remove/disable the print in production)
+ # debug print
+ # print(f"[rotary] after align: t.shape={tuple(t.shape)}, cos_.shape={tuple(cos_.shape)}")
+
+ # Final safety check: ensure cos_ can broadcast over t along seq_axis
+ # Map cos_ seq dim (we assumed it's axis 0) to the seq_axis in t via unsqueezing above.
+ # If the seq axis in t is not axis 0, the unsqueeze logic handled typical 3-D case above.
+
+ # Debug print (optional; remove when satisfied)
+ # print("rotary debug: t.shape", tuple(t.shape), "cos.shape", tuple(cos_.shape), "sin.shape", tuple(sin_.shape))
+
+ # Apply rotary
t = (t * cos_) + (_rotate_half(t) * sin_)
- return t.to(orig_dtype)
+
+ result = t.to(orig_dtype)
+
+ # DEBUG: Check output
+ if result is None:
+ print(f"[apply_rotary_pos_emb_vision] ERROR: result is None after conversion!", flush=True)
+
+ return result
class PatchEmbed(torch.nn.Module):
diff --git a/aiak_training_llm/models/mobilellm/__init__.py b/aiak_training_llm/models/mobilellm/__init__.py
new file mode 100644
index 00000000..8b78ca67
--- /dev/null
+++ b/aiak_training_llm/models/mobilellm/__init__.py
@@ -0,0 +1,13 @@
+"""MobileLLM model family for efficient vision-language models."""
+
+from .mobilellm_config import get_mobilellm_config
+from .mobilellm_model import MobileLLMModel
+from .mobilellm_layer_spec import get_mobilellm_layer_with_te_spec
+from .mobilellm_provider import mobilellm_model_provider
+
+__all__ = [
+ "get_mobilellm_config",
+ "MobileLLMModel",
+ "get_mobilellm_layer_with_te_spec",
+ "mobilellm_model_provider",
+]
diff --git a/aiak_training_llm/models/mobilellm/hf_checkpoint/.gitattributes b/aiak_training_llm/models/mobilellm/hf_checkpoint/.gitattributes
new file mode 100644
index 00000000..52373fe2
--- /dev/null
+++ b/aiak_training_llm/models/mobilellm/hf_checkpoint/.gitattributes
@@ -0,0 +1,36 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text
diff --git a/aiak_training_llm/models/mobilellm/hf_checkpoint/LICENSE b/aiak_training_llm/models/mobilellm/hf_checkpoint/LICENSE
new file mode 100644
index 00000000..65c038ca
--- /dev/null
+++ b/aiak_training_llm/models/mobilellm/hf_checkpoint/LICENSE
@@ -0,0 +1,124 @@
+FAIR Noncommercial Research License
+v1 Last Updated: August 18, 2025
+
+โAcceptable Use Policyโ means the FAIR Acceptable Use Policy, applicable to Research Materials, that is incorporated into this Agreement.
+
+โAgreementโ means the terms and conditions for use, reproduction, distribution and modification of the Research Materials set forth herein.
+
+
+โDocumentationโ means the specifications, manuals and documentation accompanying
+Research Materials distributed by Meta.
+
+
+โLicenseeโ or โyouโ means you, or your employer or any other person or entity (if you are entering into this Agreement on such person or entityโs behalf), of the age required under applicable laws, rules or regulations to provide legal consent and that has legal authority to bind your employer or such other person or entity if you are entering in this Agreement on their behalf.
+
+
+โMetaโ or โweโ means Meta Platforms Ireland Limited (if you are located in or, if you are an entity, your principal place of business is in the EEA or Switzerland) and Meta Platforms, Inc. (if you are located outside of the EEA or Switzerland).
+
+โNoncommercial Research Usesโ means noncommercial research use cases related to research, development, education, processing, or analysis and in each case, is not primarily intended for commercial advantage or monetary compensation to you or others.
+
+โResearch Materialsโ means, collectively, Documentation and the models, software and algorithms, including machine-learning model code, trained model weights, inference-enabling code, training-enabling code, fine-tuning enabling code, demonstration materials and other elements of the foregoing distributed by Meta and made available under this Agreement.
+
+By clicking โI Acceptโ below or by using or distributing any portion or element of the Research Materials, you agree to be bound by this Agreement.
+
+
+1. License Rights and Redistribution.
+
+
+a. Grant of Rights. You are granted a non-exclusive, worldwide, non-transferable and royalty-free limited license under Metaโs intellectual property or other rights owned by Meta embodied in the Research Materials to use, reproduce, distribute, copy, create derivative works of, and make modifications to the Research Materials.
+
+b. Redistribution and Use.
+ i. You will not use the Research Materials or any outputs or results of the Research Materials in connection with any commercial uses or for any uses other than Noncommercial Research Uses;
+
+
+ii. Distribution of Research Materials, and any derivative works thereof, are subject to the terms of this Agreement. If you distribute or make the Research Materials, or any derivative works thereof, available to a third party, you may only do so under the terms of this Agreement. You shall also provide a copy of this Agreement to such third party.
+
+
+iii. If you submit for publication the results of research you perform on, using, or otherwise in connection with Research Materials, you must acknowledge the use of Research Materials in your publication.
+
+
+iv. Your use of the Research Materials must comply with applicable laws and regulations (including Trade Control Laws) and adhere to the FAIR Acceptable Use Policy, which is hereby incorporated by reference into this Agreement.
+2. User Support. Your Noncommercial Research Use of the Research Materials is done at your own discretion; Meta does not process any information nor provide any service in relation to such use. Meta is under no obligation to provide any support services for the Research Materials. Any support provided is โas isโ, โwith all faultsโ, and without warranty of any kind.
+
+
+3. Disclaimer of Warranty. UNLESS REQUIRED BY APPLICABLE LAW, THE RESEARCH MATERIALS AND ANY OUTPUT AND RESULTS THEREFROM ARE PROVIDED ON AN โAS ISโ BASIS, WITHOUT WARRANTIES OF ANY KIND, AND META DISCLAIMS ALL WARRANTIES OF ANY KIND, BOTH EXPRESS AND IMPLIED, INCLUDING, WITHOUT LIMITATION, ANY WARRANTIES OF TITLE, NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. YOU ARE SOLELY RESPONSIBLE FOR DETERMINING THE APPROPRIATENESS OF USING OR REDISTRIBUTING THE RESEARCH MATERIALS AND ASSUME ANY RISKS ASSOCIATED WITH YOUR USE OF THE RESEARCH MATERIALS AND ANY OUTPUT AND RESULTS.
+
+4. Limitation of Liability. IN NO EVENT WILL META OR ITS AFFILIATES BE LIABLE UNDER ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, TORT, NEGLIGENCE, PRODUCTS LIABILITY, OR OTHERWISE, ARISING OUT OF THIS AGREEMENT, FOR ANY LOST PROFITS OR ANY DIRECT OR INDIRECT, SPECIAL, CONSEQUENTIAL, INCIDENTAL, EXEMPLARY OR PUNITIVE DAMAGES, EVEN IF META OR ITS AFFILIATES HAVE BEEN ADVISED OF THE POSSIBILITY OF ANY OF THE FOREGOING.
+
+5. Intellectual Property.
+
+
+a. Subject to Metaโs ownership of Research Materials and derivatives made by or for Meta, with respect to any derivative works and modifications of the Research Materials that are made by you, as between you and Meta, you are and will be the owner of such derivative works and modifications.
+
+b. If you institute litigation or other proceedings against Meta or any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the Research Materials, outputs or results, or any portion of any of the foregoing, constitutes infringement of intellectual property or other rights owned or licensable by you, then any licenses granted to you under this Agreement shall terminate as of the date such litigation or claim is filed or instituted. You will indemnify and hold harmless Meta from and against any claim by any third party arising out of or related to your use or distribution of the Research Materials.
+
+6. Term and Termination. The term of this Agreement will commence upon your acceptance of this Agreement or access to the Research Materials and will continue in full force and effect until terminated in accordance with the terms and conditions herein. Meta may terminate this Agreement if you are in breach of any term or condition of this Agreement. Upon termination of this Agreement, you shall delete and cease use of the Research Materials. Sections 3, 4 and 7 shall survive the termination of this Agreement.
+
+7. Governing Law and Jurisdiction. This Agreement will be governed and construed under the laws of the State of California without regard to choice of law principles, and the UN Convention on Contracts for the International Sale of Goods does not apply to this Agreement. The courts of California shall have exclusive jurisdiction of any dispute arising out of this Agreement.
+
+
+8. Modifications and Amendments. Meta may modify this Agreement from time to time; provided that they are similar in spirit to the current version of the Agreement, but may differ in detail to address new problems or concerns. All such changes will be effective immediately. Your continued use of the Research Materials after any modification to this Agreement constitutes your agreement to such modification. Except as provided in this Agreement, no modification or addition to any provision of this Agreement will be binding unless it is in writing and signed by an authorized representative of both you and Meta.
+
+
+FAIR Acceptable Use Policy
+
+The Fundamental AI Research (FAIR) team at Meta seeks to further understanding of new and existing research domains with the mission of advancing the state-of-the-art in artificial intelligence through open research for the benefit of all.
+
+As part of this mission, Meta makes certain research materials available for noncommercial research use. Meta is committed to promoting the safe and responsible use of such research materials.
+
+Prohibited Uses
+
+You agree you will not use, or allow others to use, Research Materials to:
+
+ Violate the law or othersโ rights, including to:
+Engage in, promote, generate, contribute to, encourage, plan, incite, or further illegal or unlawful activity or content, such as:
+Violence or terrorism
+Exploitation or harm to children, including the solicitation, creation, acquisition, or dissemination of child exploitative content or failure to report Child Sexual Abuse Material
+Human trafficking, exploitation, and sexual violence
+The illegal distribution of information or materials to minors, including obscene materials, or failure to employ legally required age-gating in connection with such information or materials.
+Sexual solicitation
+Any other criminal activity
+
+Engage in, promote, incite, or facilitate the harassment, abuse, threatening, or bullying of individuals or groups of individuals
+
+Engage in, promote, incite, or facilitate discrimination or other unlawful or harmful conduct in the provision of employment, employment benefits, credit, housing, other economic benefits, or other essential goods and services
+
+Engage in the unauthorized or unlicensed practice of any profession including, but not limited to, financial, legal, medical/health, or related professional practices
+
+Collect, process, disclose, generate, or infer health, demographic, or other sensitive personal or private information about individuals without rights and consents required by applicable laws
+
+Engage in or facilitate any action or generate any content that infringes, misappropriates, or otherwise violates any third-party rights, including the outputs or results of any technology using FAIR research materials
+
+Create, generate, or facilitate the creation of malicious code, malware, computer viruses or do anything else that could disable, overburden, interfere with or impair the proper working, integrity, operation or appearance of a website or computer system
+
+2. Engage in, promote, incite, facilitate, or assist in the planning or development of activities that present a risk of death or bodily harm to individuals, including use of research artifacts related to the following:
+
+Military, warfare, nuclear industries or applications, espionage, use for materials or activities that are subject to the International Traffic Arms Regulations (ITAR) maintained by the United States Department of State
+
+Guns and illegal weapons (including weapon development)
+
+Illegal drugs and regulated/controlled substances
+
+Operation of critical infrastructure, transportation technologies, or heavy machinery
+
+Self-harm or harm to others, including suicide, cutting, and eating disorders
+
+Any content intended to incite or promote violence, abuse, or any infliction of bodily harm to an individual
+
+3. Intentionally deceive or mislead others, including use of FAIR Research Materials related to the following:
+
+ Generating, promoting, or furthering fraud or the creation or promotion of disinformation
+
+ Generating, promoting, or furthering defamatory content, including the creation of defamatory statements, images, or other content
+
+Generating, promoting, or further distributing spam
+
+ Impersonating another individual without consent, authorization, or legal right
+
+Representing that outputs of FAIR research materials or outputs from technology using FAIR research materials are human-generated
+
+Generating or facilitating false online engagement, including fake reviews and other means of fake online engagement
+
+4. Fail to appropriately disclose to end users any known dangers of your Research Materials.
+
+Please report any violation of this Policy or other problems that could lead to a violation of this Policy by submitting a report here [https://docs.google.com/forms/d/e/1FAIpQLSeb11cryAopJ7LNrC4nxEUXrHY26hfkXQMf_uH-oFgA3WlYZQ/viewform].
diff --git a/aiak_training_llm/models/mobilellm/hf_checkpoint/README.md b/aiak_training_llm/models/mobilellm/hf_checkpoint/README.md
new file mode 100644
index 00000000..ba02dfb9
--- /dev/null
+++ b/aiak_training_llm/models/mobilellm/hf_checkpoint/README.md
@@ -0,0 +1,622 @@
+---
+license: other
+license_name: fair-noncommercial-research
+extra_gated_prompt: >
+ FAIR Noncommercial Research License v1 Last Updated: August 18, 2025
+
+ โAcceptable Use Policyโ means the FAIR Acceptable Use Policy, applicable to
+ Research Materials, that is incorporated into this Agreement.
+
+ โAgreementโ means the terms and conditions for use, reproduction, distribution
+ and modification of the Research Materials set forth herein.
+
+
+ โDocumentationโ means the specifications, manuals and documentation
+ accompanying Research Materials distributed by Meta.
+
+
+ โLicenseeโ or โyouโ means you, or your employer or any other person or entity
+ (if you are entering into this Agreement on such person or entityโs behalf),
+ of the age required under applicable laws, rules or regulations to provide
+ legal consent and that has legal authority to bind your employer or such other
+ person or entity if you are entering in this Agreement on their behalf.
+
+
+ โMetaโ or โweโ means Meta Platforms Ireland Limited (if you are located in or,
+ if you are an entity, your principal place of business is in the EEA or
+ Switzerland) and Meta Platforms, Inc. (if you are located outside of the EEA
+ or Switzerland).
+
+ โNoncommercial Research Usesโ means noncommercial research use cases related
+ to research, development, education, processing, or analysis and in each case,
+ is not primarily intended for commercial advantage or monetary compensation to
+ you or others.
+
+ โResearch Materialsโ means, collectively, Documentation and the models,
+ software and algorithms, including machine-learning model code, trained model
+ weights, inference-enabling code, training-enabling code, fine-tuning enabling
+ code, demonstration materials and other elements of the foregoing distributed
+ by Meta and made available under this Agreement.
+
+ By clicking โI Acceptโ below or by using or distributing any portion or
+ element of the Research Materials, you agree to be bound by this Agreement.
+
+
+ 1. License Rights and Redistribution.
+
+
+ a. Grant of Rights. You are granted a non-exclusive, worldwide,
+ non-transferable and royalty-free limited license under Metaโs intellectual
+ property or other rights owned by Meta embodied in the Research Materials to
+ use, reproduce, distribute, copy, create derivative works of, and make
+ modifications to the Research Materials.
+
+ b. Redistribution and Use. i. You will not use the Research Materials or any
+ outputs or results of the Research Materials in connection with any commercial
+ uses or for any uses other than Noncommercial Research Uses;
+
+
+ ii. Distribution of Research Materials, and any derivative works thereof, are
+ subject to the terms of this Agreement. If you distribute or make the Research
+ Materials, or any derivative works thereof, available to a third party, you
+ may only do so under the terms of this Agreement. You shall also provide a
+ copy of this Agreement to such third party.
+
+
+ iii. If you submit for publication the results of research you perform on,
+ using, or otherwise in connection with Research Materials, you must
+ acknowledge the use of Research Materials in your publication.
+
+
+ iv. Your use of the Research Materials must comply with applicable laws and
+ regulations (including Trade Control Laws) and adhere to the FAIR Acceptable
+ Use Policy, which is hereby incorporated by reference into this Agreement. 2.
+ User Support. Your Noncommercial Research Use of the Research Materials is
+ done at your own discretion; Meta does not process any information nor provide
+ any service in relation to such use. Meta is under no obligation to provide
+ any support services for the Research Materials. Any support provided is โas
+ isโ, โwith all faultsโ, and without warranty of any kind.
+
+
+ 3. Disclaimer of Warranty. UNLESS REQUIRED BY APPLICABLE LAW, THE RESEARCH
+ MATERIALS AND ANY OUTPUT AND RESULTS THEREFROM ARE PROVIDED ON AN โAS ISโ
+ BASIS, WITHOUT WARRANTIES OF ANY KIND, AND META DISCLAIMS ALL WARRANTIES OF
+ ANY KIND, BOTH EXPRESS AND IMPLIED, INCLUDING, WITHOUT LIMITATION, ANY
+ WARRANTIES OF TITLE, NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A
+ PARTICULAR PURPOSE. YOU ARE SOLELY RESPONSIBLE FOR DETERMINING THE
+ APPROPRIATENESS OF USING OR REDISTRIBUTING THE RESEARCH MATERIALS AND ASSUME
+ ANY RISKS ASSOCIATED WITH YOUR USE OF THE RESEARCH MATERIALS AND ANY OUTPUT
+ AND RESULTS.
+
+ 4. Limitation of Liability. IN NO EVENT WILL META OR ITS AFFILIATES BE LIABLE
+ UNDER ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, TORT, NEGLIGENCE, PRODUCTS
+ LIABILITY, OR OTHERWISE, ARISING OUT OF THIS AGREEMENT, FOR ANY LOST PROFITS
+ OR ANY DIRECT OR INDIRECT, SPECIAL, CONSEQUENTIAL, INCIDENTAL, EXEMPLARY OR
+ PUNITIVE DAMAGES, EVEN IF META OR ITS AFFILIATES HAVE BEEN ADVISED OF THE
+ POSSIBILITY OF ANY OF THE FOREGOING.
+
+ 5. Intellectual Property.
+
+
+ a. Subject to Metaโs ownership of Research Materials and derivatives made by
+ or for Meta, with respect to any derivative works and modifications of the
+ Research Materials that are made by you, as between you and Meta, you are and
+ will be the owner of such derivative works and modifications.
+
+ b. If you institute litigation or other proceedings against Meta or any entity
+ (including a cross-claim or counterclaim in a lawsuit) alleging that the
+ Research Materials, outputs or results, or any portion of any of the
+ foregoing, constitutes infringement of intellectual property or other rights
+ owned or licensable by you, then any licenses granted to you under this
+ Agreement shall terminate as of the date such litigation or claim is filed or
+ instituted. You will indemnify and hold harmless Meta from and against any
+ claim by any third party arising out of or related to your use or distribution
+ of the Research Materials.
+
+ 6. Term and Termination. The term of this Agreement will commence upon your
+ acceptance of this Agreement or access to the Research Materials and will
+ continue in full force and effect until terminated in accordance with the
+ terms and conditions herein. Meta may terminate this Agreement if you are in
+ breach of any term or condition of this Agreement. Upon termination of this
+ Agreement, you shall delete and cease use of the Research Materials. Sections
+ 3, 4 and 7 shall survive the termination of this Agreement.
+
+ 7. Governing Law and Jurisdiction. This Agreement will be governed and
+ construed under the laws of the State of California without regard to choice
+ of law principles, and the UN Convention on Contracts for the International
+ Sale of Goods does not apply to this Agreement. The courts of California shall
+ have exclusive jurisdiction of any dispute arising out of this Agreement.
+
+
+ 8. Modifications and Amendments. Meta may modify this Agreement from time to
+ time; provided that they are similar in spirit to the current version of the
+ Agreement, but may differ in detail to address new problems or concerns. All
+ such changes will be effective immediately. Your continued use of the Research
+ Materials after any modification to this Agreement constitutes your agreement
+ to such modification. Except as provided in this Agreement, no modification or
+ addition to any provision of this Agreement will be binding unless it is in
+ writing and signed by an authorized representative of both you and Meta.
+
+
+ FAIR Acceptable Use Policy
+
+ The Fundamental AI Research (FAIR) team at Meta seeks to further understanding
+ of new and existing research domains with the mission of advancing the
+ state-of-the-art in artificial intelligence through open research for the
+ benefit of all.
+
+ As part of this mission, Meta makes certain research materials available for
+ noncommercial research use. Meta is committed to promoting the safe and
+ responsible use of such research materials.
+
+ Prohibited Uses
+
+ You agree you will not use, or allow others to use, Research Materials to:
+
+ Violate the law or othersโ rights, including to: Engage in, promote, generate,
+ contribute to, encourage, plan, incite, or further illegal or unlawful
+ activity or content, such as: Violence or terrorism Exploitation or harm to
+ children, including the solicitation, creation, acquisition, or dissemination
+ of child exploitative content or failure to report Child Sexual Abuse Material
+ Human trafficking, exploitation, and sexual violence The illegal distribution
+ of information or materials to minors, including obscene materials, or failure
+ to employ legally required age-gating in connection with such information or
+ materials. Sexual solicitation Any other criminal activity
+
+ Engage in, promote, incite, or facilitate the harassment, abuse, threatening,
+ or bullying of individuals or groups of individuals
+
+ Engage in, promote, incite, or facilitate discrimination or other unlawful or
+ harmful conduct in the provision of employment, employment benefits, credit,
+ housing, other economic benefits, or other essential goods and services
+
+ Engage in the unauthorized or unlicensed practice of any profession including,
+ but not limited to, financial, legal, medical/health, or related professional
+ practices
+
+ Collect, process, disclose, generate, or infer health, demographic, or other
+ sensitive personal or private information about individuals without rights and
+ consents required by applicable laws
+
+ Engage in or facilitate any action or generate any content that infringes,
+ misappropriates, or otherwise violates any third-party rights, including the
+ outputs or results of any technology using FAIR research materials
+
+ Create, generate, or facilitate the creation of malicious code, malware,
+ computer viruses or do anything else that could disable, overburden, interfere
+ with or impair the proper working, integrity, operation or appearance of a
+ website or computer system
+
+ 2. Engage in, promote, incite, facilitate, or assist in the planning or
+ development of activities that present a risk of death or bodily harm to
+ individuals, including use of research artifacts related to the following:
+
+ Military, warfare, nuclear industries or applications, espionage, use for
+ materials or activities that are subject to the International Traffic Arms
+ Regulations (ITAR) maintained by the United States Department of State
+
+ Guns and illegal weapons (including weapon development)
+
+ Illegal drugs and regulated/controlled substances
+
+ Operation of critical infrastructure, transportation technologies, or heavy
+ machinery
+
+ Self-harm or harm to others, including suicide, cutting, and eating disorders
+
+ Any content intended to incite or promote violence, abuse, or any infliction
+ of bodily harm to an individual
+
+ 3. Intentionally deceive or mislead others, including use of FAIR Research
+ Materials related to the following:
+
+ Generating, promoting, or furthering fraud or the creation or promotion of
+ disinformation
+
+ Generating, promoting, or furthering defamatory content, including the
+ creation of defamatory statements, images, or other content
+
+ Generating, promoting, or further distributing spam
+
+ Impersonating another individual without consent, authorization, or legal
+ right
+
+ Representing that outputs of FAIR research materials or outputs from
+ technology using FAIR research materials are human-generated
+
+ Generating or facilitating false online engagement, including fake reviews and
+ other means of fake online engagement
+
+ 4. Fail to appropriately disclose to end users any known dangers of your
+ Research Materials.
+
+ Please report any violation of this Policy or other problems that could lead
+ to a violation of this Policy by submitting a report here
+ [https://docs.google.com/forms/d/e/1FAIpQLSeb11cryAopJ7LNrC4nxEUXrHY26hfkXQMf_uH-oFgA3WlYZQ/viewform].
+extra_gated_fields:
+ First Name: text
+ Last Name: text
+ Date of birth: date_picker
+ Country: country
+ Affiliation: text
+ Job title:
+ type: select
+ options:
+ - Student
+ - Research Graduate
+ - AI researcher
+ - AI developer/engineer
+ - Reporter
+ - Other
+ geo: ip_location
+ By clicking Submit below I accept the terms of the license and acknowledge that the information I provide will be collected stored processed and shared in accordance with the Meta Privacy Policy: checkbox
+extra_gated_description: >-
+ The information you provide will be collected, stored, processed and shared in
+ accordance with the [Meta Privacy
+ Policy](https://www.facebook.com/privacy/policy/).
+extra_gated_button_content: Submit
+extra_gated_heading: >-
+ Please be sure to provide your full legal name, date of birth, and full
+ organization name with all corporate identifiers. Avoid the use of acronyms
+ and special characters. Failure to follow these instructions may prevent you
+ from accessing this model and others on Hugging Face. You will not have the
+ ability to edit this form after submission, so please ensure all information
+ is accurate.
+language:
+- en
+library_name: transformers
+tags:
+- facebook
+- meta
+- pytorch
+- mobilellm
+base_model:
+- facebook/MobileLLM-R1-140M-base
+---
+
+
+
+
+ ๐ค Hugging Face | ๐ Paper | ๐ป Code
+
+
+# Model Details
+
+We present MobileLLM-R1, a new series of efficient reasoning models in the MobileLLM family. The release includes two categories of models:
+
+Base models:
+- [MobileLLM-R1-140M-base](https://huggingface.co/facebook/MobileLLM-R1-140M-base/)
+- [MobileLLM-R1-360M-base](https://huggingface.co/facebook/MobileLLM-R1-360M-base/)
+- [MobileLLM-R1-950M-base](https://huggingface.co/facebook/MobileLLM-R1-950M-base/)
+
+Final models:
+- [MobileLLM-R1-140M](https://huggingface.co/facebook/MobileLLM-R1-140M/)
+- [MobileLLM-R1-360M](https://huggingface.co/facebook/MobileLLM-R1-360M/)
+- [MobileLLM-R1-950M](https://huggingface.co/facebook/MobileLLM-R1-950M/)
+
+> **Note**: These models are not general-purpose chat models. They are Supervised Fine-Tuned (SFT) models, specifically trained to address mathematical, programming (Python, C++), and scientific problems.
+
+In addition to the models, we release the complete training recipes and data sources to ensure reproducibility and support further research.
+
+Remarkably, the MobileLLM-R1 950M, pre-trained on only **~2T high-quality tokens** and with fewer than 5T total training tokens, achieves comparable or superior performance to Qwen3 0.6B, which was trained on 36T tokens, across MATH, GSM8K, MMLU, and LiveCodeBench benchmarks.
+
+Compared to existing fully open-source models, MobileLLM-R1 950M model achieves **~5ร higher accuracy on MATH** compared to the Olmo 1.24B model and **~2ร higher accuracy** relative to the SmolLM2 1.7B model, despite being substantially smaller in parameter scale. In addition, MobileLLM-R1 950M outperforms both Olmo 1.24B and SmolLM2 1.7B **by a wide margin on coding benchmarks**, establishing a new state-of-the-art among fully open-source models.
+
+## News
+- Sept 12, 2025: ๐ MobileLLM-R1 models are released on [HuggingFace](https://huggingface.co/collections/facebook/mobilellm-r1-68c4597b104fac45f28f448e).
+- Spet 28, 2025: ๐ The training code is also available on [GitHub](https://github.com/facebookresearch/MobileLLM-R1).
+- Spet 29, 2025: ๐ The technical report "[MobileLLM-R1: Exploring the Limits of Sub-Billion Language Model Reasoners with Open Training Recipes](https://arxiv.org/pdf/2509.24945)" is also available! Please check it out.
+
+# Highlights
+
+
+### Pretrained Model
+
+
+### Token efficiency comparison across pretrained models
+
+
+### Post-trained Model
+
+
+**Model Architecture**:
+
+| | # Layers | # Attnetion Heads | # KV Heads | Dim | Hidden Dim | Params |
+| --- | --- | --- | --- | --- | --- | --- |
+| MobileLLM-R1-140M | 15 | 9 | 3 | 576 | 2048 | 140M |
+| MobileLLM-R1-360M | 15 | 16 | 4 | 1024 | 4096 | 359M |
+| MobileLLM-R1-950M | 22 | 24 | 6 | 1536 | 6144 | 949M |
+
+| | Input modalities | Output modalities | Context Length | Vocaburary Size | Shared Embeddings |
+| --- | --- | --- | --- | --- | --- |
+| [MobileLLM-R1-140M-base](https://huggingface.co/facebook/MobileLLM-R1-140M-base) | Text | Text | 4k | 128k | Yes |
+| [MobileLLM-R1-360M-base](https://huggingface.co/facebook/MobileLLM-R1-360M-base) | Text | Text | 4k | 128k | Yes |
+| [MobileLLM-R1-950M-base](https://huggingface.co/facebook/MobileLLM-R1-950M-base) | Text | Text | 4k | 128k | Yes |
+| [MobileLLM-R1-140M](https://huggingface.co/facebook/MobileLLM-R1-140M) | Text | Text | 32k | 128k | Yes |
+| [MobileLLM-R1-360M](https://huggingface.co/facebook/MobileLLM-R1-360M) | Text | Text | 32k | 128k | Yes |
+| [MobileLLM-R1-950M](https://huggingface.co/facebook/MobileLLM-R1-950M) | Text | Text | 32k | 128k | Yes |
+
+# How to use
+
+To load the pretrained model for further finetuning or evaluation:
+```bash
+from transformers import AutoModelForCausalLM, AutoTokenizer
+tokenizer = AutoTokenizer.from_pretrained("facebook/MobileLLM-R1-950M")
+model = AutoModelForCausalLM.from_pretrained("facebook/MobileLLM-R1-950M")
+```
+
+# Inference examples
+
+Transformers
+
+```py
+from transformers import pipeline
+import torch
+
+model_id = "facebook/MobileLLM-R1-950M"
+
+pipe = pipeline(
+ "text-generation",
+ model=model_id,
+ torch_dtype="auto",
+ device_map="auto",
+)
+
+# Math problem / default scenario
+messages = [
+ {
+ "role": "system",
+ "content": "Please reason step by step, and put your final answer within \\boxed{}."
+ },
+ {"role": "user", "content": "Compute: $1-2+3-4+5- \\dots +99-100$."},
+]
+
+# C++ coding scenario
+messages = [
+ {
+ "role": "system",
+ "content": (
+ "\nYou are a helpful and harmless assistant. You should think step-by-step before responding to the instruction below.\n\n"
+ "Please use c++ programming language only.\n"
+ "You must use ```cpp for just the final solution code block with the following format:\n"
+ "```cpp\n# Your code here\n```\n"
+ )
+ },
+ {"role": "user", "content": "Write a C++ program that prints 'Hello, World!'."},
+]
+
+# Python coding scenario
+messages = [
+ {
+ "role": "system",
+ "content": (
+ "\nYou are a helpful and harmless assistant. You should think step-by-step before responding to the instruction below.\n\n"
+ "Please use python programming language only.\n"
+ "You must use ```python for just the final solution code block with the following format:\n"
+ "```python\n# Your code here\n```\n"
+ )
+ },
+ {"role": "user", "content": "Write a Python function that returns the square of a number."},
+]
+
+outputs = pipe(
+ messages,
+ max_new_tokens=8192,
+)
+print(outputs[0]["generated_text"][-1])
+```
+
+You can also run inference with vLLM. You only need to register the model architecture Llama4ForCausalLM with the vLLM ModelRegistry.
+```bash
+from vllm.model_executor.models.llama4 import Llama4ForCausalLM
+from vllm.model_executor.models.registry import ModelRegistry
+ModelRegistry.register_model("Llama4ForCausalLM", Llama4ForCausalLM)
+```
+
+
+# Evaluation
+
+## MobileLLM-R1 base model
+| Model | Size | MATH500 | GSM8K | MBPP | HumanEval | CommonSense Avg. | MMLU |
+| --- | --- | --- | --- | --- | --- | --- | --- |
+| | | 4-shot
+ print("Starting training...")
# start training
trainer.train()
diff --git a/aiak_training_llm/train/arguments.py b/aiak_training_llm/train/arguments.py
index acdf5b24..fc57272e 100644
--- a/aiak_training_llm/train/arguments.py
+++ b/aiak_training_llm/train/arguments.py
@@ -79,7 +79,7 @@ def _add_extra_training_rice_vl_args(parser: argparse.ArgumentParser) -> argpars
group.add_argument(
'--training-rice-vl-max-answer-length',
type=int,
- default=4096,
+ default=512, # Changed from 4096 to 512 to match seq_length
help=(
"The maximum number of characters allowed in an answer during training. "
"Answers longer than this will be truncated."
@@ -415,6 +415,22 @@ def _add_extra_multimodal_args(parser):
group.add_argument('--fps-max-frames', type=int, default=768,
help='The maximum number of frames of the video')
+
+ # FastViT specific arguments (following FastVLM repo)
+ group.add_argument('--use-fastvit', action='store_true', default=False,
+ help='Use FastViT vision encoder instead of Rice/SigLIP')
+ group.add_argument('--fastvit-image-size', type=int, default=384,
+ help='FastViT input image size (default: 384)')
+ group.add_argument('--vision-tower-name', type=str, default='mobileclip_l_384',
+ help='FastViT model variant (mobileclip_l_384, mobileclip_l_448, mobileclip_l_512)')
+ group.add_argument('--image-aspect-ratio', type=str, default='pad',
+ choices=['pad', 'anyres', 'square'],
+ help='Image aspect ratio handling: pad (expand to square with padding), '
+ 'anyres (variable resolution with patches), square (direct resize)')
+ group.add_argument('--image-grid-pinpoints', type=str,
+ default='[(384, 384), (768, 384), (384, 768), (768, 768)]',
+ help='Grid pinpoints for anyres image processing')
+
return parser
@@ -436,6 +452,8 @@ def _validate_extra_model_args(args):
if model_config is not None:
# the structural configuration of model will be overwritten, such as num_layers, hidden_states..
print_rank_0(f'-------------- Configure model to {args.model_name} --------------', args.rank)
+ print_rank_0(f'[DEBUG] Model config type: {type(model_config).__name__}', args.rank)
+ print_rank_0(f'[DEBUG] Is MobileLLM: {"mobilellm" in args.model_name.lower()}', args.rank)
for field in fields(model_config.__class__):
assert hasattr(args, field.name), f"The model config field ({field.name}) is not defined in args."
@@ -444,6 +462,8 @@ def _validate_extra_model_args(args):
print_rank_0(f" {key} = {value} ", args.rank)
print_rank_0('---------------- End of configuration ----------------', args.rank)
+ print_rank_0(f'[DEBUG] Args now has: num_layers={args.num_layers}, hidden_size={args.hidden_size}, '
+ f'num_attention_heads={args.num_attention_heads}, vocab_size={args.vocab_size_in_config_file}', args.rank)
if args.enable_fa_within_mla:
args.attention_backend = AttnBackend.flash
diff --git a/aiak_training_llm/train/megatron_trainer.py b/aiak_training_llm/train/megatron_trainer.py
index b98b62cd..f23f8608 100644
--- a/aiak_training_llm/train/megatron_trainer.py
+++ b/aiak_training_llm/train/megatron_trainer.py
@@ -43,15 +43,17 @@ def __init__(
def train(self):
"""start training"""
+ # Mark dataset provider as distributed (multi-GPU)
self.train_valid_test_datasets_provider.is_distributed = True
pretrain(
- train_args=self.train_args,
+ train_args=self.train_args, #All training arguments (like batch size, lr, etc)
train_valid_test_dataset_provider=self.train_valid_test_datasets_provider,
- model_provider=self.model_provider,
- model_type=self.model_type,
- forward_step_func=self.forward_step_func,
- process_non_loss_data_func=self.process_non_loss_data_func,
- get_embedding_ranks=self.get_embedding_ranks,
- get_position_embedding_ranks=self.get_position_embedding_ranks,
+ # function that provides train/valid/test datasets (preprocessing handled inside)
+ model_provider=self.model_provider, #function that provides the model architecture
+ model_type=self.model_type, #type of model being trained (e.g., language model, vision model, etc)(encoder or decoder)
+ forward_step_func=self.forward_step_func, #function that computes loss for one batch
+ process_non_loss_data_func=self.process_non_loss_data_func, #Optional function for logging/visualization
+ get_embedding_ranks=self.get_embedding_ranks, #for distributed embedding tables
+ get_position_embedding_ranks=self.get_position_embedding_ranks, # for distributed position embeddings
)
diff --git a/aiak_training_llm/train/pretrain/pretrain_llavaov_1_5.py b/aiak_training_llm/train/pretrain/pretrain_llavaov_1_5.py
index 4f065894..e53d750d 100644
--- a/aiak_training_llm/train/pretrain/pretrain_llavaov_1_5.py
+++ b/aiak_training_llm/train/pretrain/pretrain_llavaov_1_5.py
@@ -30,10 +30,30 @@
stimer = StragglerDetector()
-# TODO: get token id from tokenizer
-image_token_id = 151655
-video_token_id = 151656
-vision_start_token_id = 151652
+# Resolved lazily from the tokenizer so MobileLLM (128258/128259) and
+# Qwen2-VL (151655/151656) both work without hardcoding.
+_resolved_token_ids = None
+
+def _get_vision_token_ids():
+ """Return (image_token_id, video_token_id) from the runtime tokenizer.
+ Falls back to Qwen2-VL defaults if the tokenizer does not define these
+ special tokens. Result is cached after the first call."""
+ global _resolved_token_ids
+ if _resolved_token_ids is not None:
+ return _resolved_token_ids
+ img_id, vid_id = 151655, 151656 # Qwen2-VL fallback defaults
+ try:
+ tok = get_tokenizer()
+ _img = tok.convert_tokens_to_ids("<|image_pad|>")
+ _vid = tok.convert_tokens_to_ids("<|video_pad|>")
+ if _img is not None and _img > 0:
+ img_id = int(_img)
+ if _vid is not None and _vid > 0:
+ vid_id = int(_vid)
+ except Exception:
+ pass
+ _resolved_token_ids = (img_id, vid_id)
+ return _resolved_token_ids
def qwen2vl_embedding_ranks(pp_ranks):
@@ -58,7 +78,7 @@ def qwen2vl_position_embedding_ranks(pp_ranks):
Args:
pp_ranks: A list of global ranks that constitute a pipeline group.
"""
- args = get_args()
+ args = get_args() #get training arguments
# encoder size is also the index to the first rank of the decoder.
epp = args.encoder_pipeline_model_parallel_size or 0
@@ -81,9 +101,14 @@ def model_provider(pre_process=True, post_process=True, add_encoder=True, add_de
MCoreModel: The returned model
"""
args = get_args()
+ # Get model family: "llava-ov-1.5-4b" โ "llava_ov_1_5"
model_family = get_model_family(args.model_name)
+ # Lookup the registered provider
model_provider = get_model_provider(model_family)
+ # = MODEL_FAMILY_TO_PROVIDER["llava_ov_1_5"]
+ # = rice_vl_model_provider
assert model_provider is not None, f'model provider for {args.model_name} not found'
+ # call the provider to build the model and return the model
return model_provider(pre_process, post_process, add_encoder, add_decoder)
@@ -113,6 +138,7 @@ def get_batch(data_iterator):
cu_lengths = tensor_parallel.broadcast_data(["cu_lengths"], data, torch.int32)["cu_lengths"]
max_lengths = tensor_parallel.broadcast_data(["max_lengths"], data, torch.int32)["max_lengths"]
+ image_token_id, video_token_id = _get_vision_token_ids()
has_video = video_token_id in tokens
has_image = image_token_id in tokens
thw = None
@@ -235,7 +261,7 @@ def loss_func(loss_mask: torch.Tensor, output_tensor: torch.Tensor):
loss_reduced_dict
)
-
+#called by megatron training loop during each training step (called once per micro batch during training)
def forward_step(data_iterator, model):
"""Forward training step.
@@ -250,14 +276,28 @@ def forward_step(data_iterator, model):
global stimer
with stimer(bdata=True):
+ #get batch from data iterator
images, image_grid_thw, pixel_values_videos, video_grid_thw, \
input_ids, position_ids, attention_mask, \
labels, loss_mask, attn_mask_type, packed_seq_params \
= get_batch(data_iterator)
+ #returns:
+ # images: Tensor([num_images, C, H, W]) # Processed image pixels
+ # image_grid_thw: Tensor([num_images, 3]) # Grid dimensions [t, h, w]
+ # pixel_values_videos: Tensor or None # Video frames (if present)
+ # video_grid_thw: Tensor or None # Video grid dimensions
+ # input_ids: Tensor([batch, seq_len]) # Token IDs (with <|image_pad|>)
+ # position_ids: Tensor([batch, seq_len]) or None # Position indices
+ # attention_mask: Tensor([batch, seq_len]) # Attention mask (False=attend, True=mask)
+ # labels: Tensor([batch, seq_len]) # Target labels for loss
+ # loss_mask: Tensor([batch, seq_len]) # Which tokens to compute loss on
+ # attn_mask_type: AttnMaskType # Causal or padding_causal
+ # packed_seq_params: PackedSeqParams or None # For packed sequences
timers('batch-generator').stop()
with stimer:
+ # MODEL FORWARD PASS
output_tensor = model(
images,
image_grid_thw,
@@ -273,16 +313,40 @@ def forward_step(data_iterator, model):
return output_tensor, partial(loss_func, loss_mask)
-
+# Factory function that creates and returns train/valid/test data loaders for the Megatron trainer.
def train_valid_test_dataset_provider(train_val_test_num_samples):
""" Provides the datasets used by the trainer """
args = get_args()
+ # create task encoder
task_encoder = Qwen2VLTaskEncoder(args)
+ # Processes images/videos using processor
+ # Applies chat template to conversations
+ # Tokenizes text
+ # Creates attention masks
+ # Handles vision token placement (<|image_pad|>, <|video_pad|>)
+ # Packs sequences efficiently
+
+ #get train dataset
train_dataset = get_train_dataset(task_encoder)
+
+ #Purpose: Combines multiple samples into a batch
+ # Input: List of individual samples (each with different sequence lengths)
+ # Output: Batched tensors with padding
collator = build_sft_data_collator(DataCollatorForSeq2Seq)
+ # Example:
+ # Sample 1: [151652, 8932, 1234, ...] # Length 50
+ # Sample 2: [151652, 9821, 5567, ...] # Length 120
+ # โ (collate)
+ # Batch: [[151652, 8932, 1234, ..., , ], # Padded to 120
+ # [151652, 9821, 5567, ..., ...., ....]] # Already 120
+
+ #create dataloader
train_dataloader = get_train_loader(train_dataset, collator)
- return train_dataloader, None, None
+
+ return train_dataloader, None, None
+ # For SFT, typically only training loop is needed
+ # valid_iterator and test_iterator are set to None
@register_model_trainer(
diff --git a/aiak_training_llm/train/pretrain/pretrain_qwen2_vl.py b/aiak_training_llm/train/pretrain/pretrain_qwen2_vl.py
index ed3c407b..42a420ca 100644
--- a/aiak_training_llm/train/pretrain/pretrain_qwen2_vl.py
+++ b/aiak_training_llm/train/pretrain/pretrain_qwen2_vl.py
@@ -331,7 +331,7 @@ def forward_step(data_iterator, model):
def train_valid_test_dataset_provider(train_val_test_num_samples):
""" Provides the datasets used by the trainer """
-
+ print("pretraining using Qwen2VLTaskEncoder data provider...")
task_encoder = Qwen2VLTaskEncoder()
train_dataset = get_train_dataset(task_encoder)
collator = build_sft_data_collator(DataCollatorForSeq2Seq)
diff --git a/aiak_training_llm/train/sft/sft_llavaov_1_5_vl.py b/aiak_training_llm/train/sft/sft_llavaov_1_5_vl.py
index 5f1cee1e..f2214069 100644
--- a/aiak_training_llm/train/sft/sft_llavaov_1_5_vl.py
+++ b/aiak_training_llm/train/sft/sft_llavaov_1_5_vl.py
@@ -67,6 +67,16 @@ def get_batch(data_iterator):
else:
data = None
+ print("=" * 80)
+ print("[DEBUG GET_BATCH] Checking data from iterator:")
+ print(f" data is None: {data is None}")
+ if data is not None:
+ print(f" data keys: {data.keys() if isinstance(data, dict) else 'not a dict'}")
+ if isinstance(data, dict) and 'images' in data:
+ print(f" data['images'] is None: {data['images'] is None}")
+ print(f" data['images'] shape: {data['images'].shape if data['images'] is not None else 'None'}")
+ print("=" * 80)
+
data_i = tensor_parallel.broadcast_data([
"input_ids",
# "position_ids",
@@ -78,6 +88,11 @@ def get_batch(data_iterator):
], data, torch.int64)
data_f = tensor_parallel.broadcast_data(["images"], data, torch.float32)
+ print("[DEBUG GET_BATCH] After broadcast:")
+ print(f" data_f['images'] is None: {data_f['images'] is None}")
+ print(f" data_f['images'] shape: {data_f['images'].shape if data_f['images'] is not None else 'None'}")
+ print("=" * 80)
+
# slice batch along sequence dimension for context parallelism
assert mpu.get_context_parallel_world_size() == 1, "not implemented"
@@ -188,7 +203,7 @@ def train_valid_test_datasets_provider(train_val_test_num_samples):
tokenizer = get_tokenizer()
- processor = AutoProcessor.from_pretrained(args.hf_tokenizer_path, trust_remote_code=True)
+ processor = AutoProcessor.from_pretrained(args.hf_tokenizer_path, trust_remote_code=True, local_files_only=False)
if args.image_resolution:
setattr(processor, "image_resolution", args.image_resolution)
@@ -249,16 +264,21 @@ def train_valid_test_datasets_provider(train_val_test_num_samples):
return train_iter, valid_iter, test_iter
-
+# This registers the function for model family "llava_ov_1_5" and phase "sft"
@register_model_trainer(model_family=[constants.VisionLanguageModelFamilies.LLAVA_OV_1_5],
training_phase=constants.TrainingPhase.SFT)
def default_pretrain_trainer(train_args):
"""build trainer"""
from aiak_training_llm.train.pretrain import pretrain_llavaov_1_5
+ #imports the module containing model/dataset/forward functions for LLaVA-OV-1.5
if train_args.encoder_pipeline_model_parallel_size in [0, None]:
- model_type = ModelType.encoder_or_decoder
+ #check if model use pipeline parallelism with separate encoder and decoder stages
+ model_type = ModelType.encoder_or_decoder # unified model
+ print_rank_0("Using unified model type for training.")
else:
model_type = ModelType.encoder_and_decoder
+ print_rank_0("Using encoder-and-decoder model type for training.")
+ #created megatron trainer instance
trainer = MegatronTrainer(
train_args=train_args,
train_valid_test_dataset_provider=pretrain_llavaov_1_5.train_valid_test_dataset_provider,
diff --git a/aiak_training_llm/train/sft/sft_qwen2_vl.py b/aiak_training_llm/train/sft/sft_qwen2_vl.py
index 8d9aefbe..02b9acec 100644
--- a/aiak_training_llm/train/sft/sft_qwen2_vl.py
+++ b/aiak_training_llm/train/sft/sft_qwen2_vl.py
@@ -188,7 +188,7 @@ def train_valid_test_datasets_provider(train_val_test_num_samples):
tokenizer = get_tokenizer()
- processor = AutoProcessor.from_pretrained(args.hf_tokenizer_path, trust_remote_code=True)
+ processor = AutoProcessor.from_pretrained(args.hf_tokenizer_path, trust_remote_code=True, local_files_only=False)
if args.image_resolution:
setattr(processor, "image_resolution", args.image_resolution)
diff --git a/aiak_training_llm/train/trainer_builder.py b/aiak_training_llm/train/trainer_builder.py
index a5af5bea..565bf15d 100644
--- a/aiak_training_llm/train/trainer_builder.py
+++ b/aiak_training_llm/train/trainer_builder.py
@@ -7,7 +7,7 @@
MODEL_FAMILY_TRAINER_FACTORY = {}
-
+# decorator factory to register model trainers
def register_model_trainer(model_family: Union[str, List[str]], training_phase: str, training_func: Callable = None):
"""
register model training function
@@ -19,23 +19,27 @@ def register_model_trainer(model_family: Union[str, List[str]], training_phase:
training_phase: need to be consistent with the --training-phase definition in train.arguments
trainig_func: training function.
"""
+ # add to the factory
def _add_trainer(families, phase, func):
- if not isinstance(families, list):
+ #convert to list if single string
+ if not isinstance(families, list):
families = [families]
+ #loop through all families
for _family in families:
- _family = _family.lower()
+ _family = _family.lower() # make it case-insensitive
+ #create entry for this family if not exist
if _family not in MODEL_FAMILY_TRAINER_FACTORY:
MODEL_FAMILY_TRAINER_FACTORY[_family] = {}
-
+ #check for duplicate registration (family+phase)
if phase in MODEL_FAMILY_TRAINER_FACTORY[_family]:
raise ValueError(f"Cannot register duplicate trainer ({_family} family, {phase} phase)")
-
+ #register the trainer function
MODEL_FAMILY_TRAINER_FACTORY[_family][phase] = func
def _register_function(fn):
_add_trainer(model_family, training_phase, fn)
- return fn
+ return fn #fn is the training function being decorated
if training_func is not None:
return _add_trainer(model_family, training_phase, training_func)
@@ -43,19 +47,32 @@ def _register_function(fn):
return _register_function
-def build_model_trainer(args):
+def build_model_trainer(args): # all args including model name and training phase
"""create model trainer"""
# get model family name
- model_family = get_model_family(args.model_name)
-
+ #args.model_name="llava-ov-1.5-4b" in our case
+ model_family = get_model_family(args.model_name) #map model name to model family
+ print(f"Model family: {model_family}")
+ #returns family name like model_family="llava_ov_1_5" in our case
# get model family trainer
+
+ #check if this model family has a registered trainer
+ # MODEL_FAMILY_TRAINER_FACTORY is a nested dictionary ={family:{phase:trainer_func}}
if model_family not in MODEL_FAMILY_TRAINER_FACTORY:
raise ValueError(f"Not found trainer for {args.model_name} (family: {model_family})")
-
+ # check if this training phase has a registered trainer for this model family
+ # args.training_phase="sft" in our case
+ print(f"Training phase: {args.training_phase}")
if args.training_phase not in MODEL_FAMILY_TRAINER_FACTORY[model_family]:
raise ValueError(f"AIAK not support {args.training_phase} phase for {args.model_name} (family: {model_family})")
+ #retrieve the trainer function for this model family and training phase from the factory
+ # MODEL_FAMILY_TRAINER_FACTORY[llava_ov_1_5][sft] -> returns the trainer function
trainer = MODEL_FAMILY_TRAINER_FACTORY[model_family][args.training_phase]
- return trainer(args)
+ # This retrieves the default_pretrain_trainer function defined in sft_llavaov_1_5_vl.py
+ return trainer(args) # Calls default_pretrain_trainer(args)
+ # returns a trainer object for the specified model and training phase
+ #trainer(args) calls the trainer function with args.
+ #trainer function creates a megatron trainer object and returns it.
diff --git a/aiak_training_llm/train/training_utils.py b/aiak_training_llm/train/training_utils.py
index e5df5ef3..bd5273cf 100644
--- a/aiak_training_llm/train/training_utils.py
+++ b/aiak_training_llm/train/training_utils.py
@@ -244,6 +244,40 @@ def pretrain(
timers('model-and-optimizer-setup').stop()
print_datetime('after model, optimizer, and learning rate scheduler are built')
config = get_model_config(model[0])
+
+ # Print full model structure
+ print_rank_0('=' * 80)
+ print_rank_0('FULL MODEL STRUCTURE:')
+ print_rank_0('=' * 80)
+ if isinstance(model, list):
+ for idx, m in enumerate(model):
+ print_rank_0(f'\n--- Model {idx} (pipeline rank {idx}) ---')
+ print_rank_0(m)
+ else:
+ print_rank_0(model)
+ print_rank_0('=' * 80)
+
+ # Print parameter trainability status
+ print_rank_0('\n' + '=' * 80)
+ print_rank_0('PARAMETER TRAINABILITY STATUS:')
+ print_rank_0('=' * 80)
+ model_to_check = model[0] if isinstance(model, list) else model
+ trainable_params = 0
+ frozen_params = 0
+ for name, param in model_to_check.named_parameters():
+ status = "TRAINABLE" if param.requires_grad else "FROZEN"
+ print_rank_0(f"{status:12s} | {name:80s} | shape: {str(tuple(param.shape)):30s} | dtype: {param.dtype}")
+ if param.requires_grad:
+ trainable_params += param.numel()
+ else:
+ frozen_params += param.numel()
+
+ print_rank_0('=' * 80)
+ print_rank_0(f'Total trainable parameters: {trainable_params:,} ({trainable_params/1e6:.2f}M)')
+ print_rank_0(f'Total frozen parameters: {frozen_params:,} ({frozen_params/1e6:.2f}M)')
+ print_rank_0(f'Total parameters: {trainable_params + frozen_params:,} ({(trainable_params + frozen_params)/1e6:.2f}M)')
+ print_rank_0(f'Trainable percentage: {100 * trainable_params / (trainable_params + frozen_params):.2f}%')
+ print_rank_0('=' * 80)
# Data stuff.
timers('train/valid/test-data-iterators-setup', log_level=0).start(barrier=True)
@@ -279,6 +313,20 @@ def pretrain(
iteration = 0
if args.do_train and args.train_iters > 0:
+ print_rank_0('\n' + '=' * 80)
+ print_rank_0("training with the following parameter status:")
+ print_rank_0('=' * 80)
+ model_to_check = model[0] if isinstance(model, list) else model
+ trainable_params = 0
+ frozen_params = 0
+ for name, param in model_to_check.named_parameters():
+ status = "TRAINABLE" if param.requires_grad else "FROZEN"
+ print_rank_0(f"{status:12s} | {name:80s} | shape: {str(tuple(param.shape)):30s} | dtype: {param.dtype}")
+ if param.requires_grad:
+ trainable_params += param.numel()
+ else:
+ frozen_params += param.numel()
+
iteration, num_floating_point_operations_so_far = train(
forward_step_func=forward_step_func,
model=model,
@@ -360,6 +408,7 @@ def setup_model_and_optimizer(model_provider_func,
timers = get_timers()
model = get_model(model_provider_func, model_type)
+ print(model)
unwrapped_model = unwrap_model(model)
kwargs = {}
@@ -403,6 +452,32 @@ def setup_model_and_optimizer(model_provider_func,
optimizer.reload_model_params()
print_rank_0(f'Upcycled checkpoint saved to {args.save}')
+ # When using FastViT with pretrained checkpoint, we need special handling
+ # The checkpoint may contain incompatible language model weights (e.g., Qwen2-1.5B vs MobileLLM-140M)
+ # So we only load the vision tower weights if pretrained_checkpoint is specified
+ if getattr(args, 'use_fastvit', False):
+ print_rank_0(f'[DEBUG] FastViT enabled: use_fastvit={getattr(args, "use_fastvit", False)}')
+ print_rank_0(f'[DEBUG] Before handling: args.load={args.load}, args.pretrained_checkpoint={args.pretrained_checkpoint}')
+
+ if args.load is not None:
+ # Keep explicit resume checkpoints. If the resume directory does not
+ # contain a checkpoint yet, Megatron can still fall back to
+ # args.pretrained_checkpoint below.
+ print_rank_0(f'FastViT enabled: Will resume/load checkpoint from: {args.load}')
+ elif args.pretrained_checkpoint is not None:
+ # We have a pretrained checkpoint - load only vision tower from HF checkpoint
+ print_rank_0(f'FastViT enabled: Will load vision tower from pretrained checkpoint: {args.pretrained_checkpoint}')
+ # Keep args.pretrained_checkpoint for vision tower loading
+ # Clear args.load to prevent Megatron checkpoint loading
+ args.load = None
+ else:
+ # No pretrained checkpoint - train vision from scratch
+ args.load = None
+ args.pretrained_checkpoint = None
+ print_rank_0('FastViT enabled: No pretrained checkpoint, training vision from scratch')
+
+ print_rank_0(f'[DEBUG] After handling: args.load={args.load}, args.pretrained_checkpoint={args.pretrained_checkpoint}')
+
if (args.load is not None or args.pretrained_checkpoint is not None) and not args.moe_use_upcycling:
timers('load-checkpoint', log_level=0).start(barrier=True)
diff --git a/aiak_training_llm/utils/constants.py b/aiak_training_llm/utils/constants.py
index c42c3394..04cacb30 100644
--- a/aiak_training_llm/utils/constants.py
+++ b/aiak_training_llm/utils/constants.py
@@ -73,6 +73,7 @@ class LanguageModelFamilies(_BaseFamilies):
QWEN3 = "qwen3"
MIXTRAL = "mixtral"
DEEPSEEK = "deepseek"
+ MOBILELLM = "mobilellm"
class VideoLanguageModelFamilies(_BaseFamilies):
diff --git a/aiak_training_llm/utils/initialize.py b/aiak_training_llm/utils/initialize.py
index 40e29448..e77fa032 100644
--- a/aiak_training_llm/utils/initialize.py
+++ b/aiak_training_llm/utils/initialize.py
@@ -46,6 +46,7 @@ def parse_arguments(
ignore_unknown_args=False,
):
"""Parse arguments."""
+ print("parse arguments function called")
args = parse_args(extra_args_provider, ignore_unknown_args)
# Prep for checkpoint conversion.
diff --git a/apex/.github/ISSUE_TEMPLATE/bug_report.md b/apex/.github/ISSUE_TEMPLATE/bug_report.md
new file mode 100644
index 00000000..63d76737
--- /dev/null
+++ b/apex/.github/ISSUE_TEMPLATE/bug_report.md
@@ -0,0 +1,23 @@
+---
+name: Bug report
+about: Create a report to help us improve apex
+title: ''
+labels: bug
+assignees: ''
+
+---
+
+**Describe the Bug**
+
+**Minimal Steps/Code to Reproduce the Bug**
+
+
+**Expected Behavior**
+
+
+**Environment**
+
diff --git a/apex/.gitignore b/apex/.gitignore
new file mode 100644
index 00000000..d30f85c3
--- /dev/null
+++ b/apex/.gitignore
@@ -0,0 +1,147 @@
+apex.egg-info
+dist
+build
+docs/build
+*~
+__pycache__
+.vscode
+
+# Copied from https://raw.githubusercontent.com/github/gitignore/master/Python.gitignore
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+# Usually these files are written by a python script from a template
+# before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+# For a library or package, you might want to ignore these files since the code is
+# intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+# However, in case of collaboration, if having platform-specific dependencies or dependencies
+# having no cross-platform support, pipenv may install dependencies that don't work, or not
+# install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
diff --git a/apex/.gitmodules b/apex/.gitmodules
new file mode 100644
index 00000000..6479428d
--- /dev/null
+++ b/apex/.gitmodules
@@ -0,0 +1,7 @@
+[submodule "apex/contrib/csrc/multihead_attn/cutlass"]
+ path = apex/contrib/csrc/multihead_attn/cutlass
+ url = https://github.com/NVIDIA/cutlass.git
+ branch = v1.2.0
+[submodule "apex/contrib/csrc/cudnn-frontend"]
+ path = apex/contrib/csrc/cudnn-frontend
+ url = https://github.com/NVIDIA/cudnn-frontend.git
diff --git a/tools/convert_checkpoint/custom/cogvlm/__init__.py b/apex/.nojekyll
similarity index 100%
rename from tools/convert_checkpoint/custom/cogvlm/__init__.py
rename to apex/.nojekyll
diff --git a/apex/LICENSE b/apex/LICENSE
new file mode 100644
index 00000000..3d1e9454
--- /dev/null
+++ b/apex/LICENSE
@@ -0,0 +1,11 @@
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
\ No newline at end of file
diff --git a/apex/README.md b/apex/README.md
new file mode 100644
index 00000000..32fdd927
--- /dev/null
+++ b/apex/README.md
@@ -0,0 +1,182 @@
+# Introduction
+
+This repository holds NVIDIA-maintained utilities to streamline mixed precision and distributed training in Pytorch.
+Some of the code here will be included in upstream Pytorch eventually.
+The intent of Apex is to make up-to-date utilities available to users as quickly as possible.
+
+## Full API Documentation: [https://nvidia.github.io/apex](https://nvidia.github.io/apex)
+
+## [GTC 2019](https://github.com/mcarilli/mixed_precision_references/tree/master/GTC_2019) and [Pytorch DevCon 2019](https://github.com/mcarilli/mixed_precision_references/tree/master/Pytorch_Devcon_2019) Slides
+
+# Contents
+
+## 1. Amp: Automatic Mixed Precision
+
+**Deprecated. Use [PyTorch AMP](https://pytorch.org/docs/stable/amp.html)**
+
+`apex.amp` is a tool to enable mixed precision training by changing only 3 lines of your script.
+Users can easily experiment with different pure and mixed precision training modes by supplying
+different flags to `amp.initialize`.
+
+[Webinar introducing Amp](https://info.nvidia.com/webinar-mixed-precision-with-pytorch-reg-page.html)
+(The flag `cast_batchnorm` has been renamed to `keep_batchnorm_fp32`).
+
+[API Documentation](https://nvidia.github.io/apex/amp.html)
+
+[Comprehensive Imagenet example](https://github.com/NVIDIA/apex/tree/master/examples/imagenet)
+
+[DCGAN example coming soon...](https://github.com/NVIDIA/apex/tree/master/examples/dcgan)
+
+[Moving to the new Amp API](https://nvidia.github.io/apex/amp.html#transition-guide-for-old-api-users) (for users of the deprecated "Amp" and "FP16_Optimizer" APIs)
+
+## 2. Distributed Training
+
+**`apex.parallel.DistributedDataParallel` is deprecated. Use [`torch.nn.parallel.DistributedDataParallel`](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html?highlight=distributeddataparallel#torch.nn.parallel.DistributedDataParallel)**
+
+`apex.parallel.DistributedDataParallel` is a module wrapper, similar to
+`torch.nn.parallel.DistributedDataParallel`. It enables convenient multiprocess distributed training,
+optimized for NVIDIA's NCCL communication library.
+
+[API Documentation](https://nvidia.github.io/apex/parallel.html)
+
+[Python Source](https://github.com/NVIDIA/apex/tree/master/apex/parallel)
+
+[Example/Walkthrough](https://github.com/NVIDIA/apex/tree/master/examples/simple/distributed)
+
+The [Imagenet example](https://github.com/NVIDIA/apex/tree/master/examples/imagenet)
+shows use of `apex.parallel.DistributedDataParallel` along with `apex.amp`.
+
+### Synchronized Batch Normalization
+
+**Deprecated. Use [`torch.nn.SyncBatchNorm`](https://pytorch.org/docs/stable/generated/torch.nn.SyncBatchNorm.html)**
+
+`apex.parallel.SyncBatchNorm` extends `torch.nn.modules.batchnorm._BatchNorm` to
+support synchronized BN.
+It allreduces stats across processes during multiprocess (DistributedDataParallel) training.
+Synchronous BN has been used in cases where only a small
+local minibatch can fit on each GPU.
+Allreduced stats increase the effective batch size for the BN layer to the
+global batch size across all processes (which, technically, is the correct
+formulation).
+Synchronous BN has been observed to improve converged accuracy in some of our research models.
+
+### Checkpointing
+
+To properly save and load your `amp` training, we introduce the `amp.state_dict()`, which contains all `loss_scalers` and their corresponding unskipped steps,
+as well as `amp.load_state_dict()` to restore these attributes.
+
+In order to get bitwise accuracy, we recommend the following workflow:
+```python
+# Initialization
+opt_level = 'O1'
+model, optimizer = amp.initialize(model, optimizer, opt_level=opt_level)
+
+# Train your model
+...
+with amp.scale_loss(loss, optimizer) as scaled_loss:
+ scaled_loss.backward()
+...
+
+# Save checkpoint
+checkpoint = {
+ 'model': model.state_dict(),
+ 'optimizer': optimizer.state_dict(),
+ 'amp': amp.state_dict()
+}
+torch.save(checkpoint, 'amp_checkpoint.pt')
+...
+
+# Restore
+model = ...
+optimizer = ...
+checkpoint = torch.load('amp_checkpoint.pt')
+
+model, optimizer = amp.initialize(model, optimizer, opt_level=opt_level)
+model.load_state_dict(checkpoint['model'])
+optimizer.load_state_dict(checkpoint['optimizer'])
+amp.load_state_dict(checkpoint['amp'])
+
+# Continue training
+...
+```
+
+Note that we recommend restoring the model using the same `opt_level`. Also note that we recommend calling the `load_state_dict` methods after `amp.initialize`.
+
+# Installation
+Each [`apex.contrib`](./apex/contrib) module requires one or more install options other than `--cpp_ext` and `--cuda_ext`.
+Note that contrib modules do not necessarily support stable PyTorch releases.
+
+## Containers
+NVIDIA PyTorch Containers are available on NGC: https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch.
+The containers come with all the custom extensions available at the moment.
+
+See [the NGC documentation](https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/index.html) for details such as:
+- how to pull a container
+- how to run a pulled container
+- release notes
+
+## From Source
+
+To install Apex from source, we recommend using the nightly Pytorch obtainable from https://github.com/pytorch/pytorch.
+
+The latest stable release obtainable from https://pytorch.org should also work.
+
+### Linux
+For performance and full functionality, we recommend installing Apex with
+CUDA and C++ extensions via
+```bash
+git clone https://github.com/NVIDIA/apex
+cd apex
+pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./
+```
+
+APEX also supports a Python-only build via
+```bash
+pip install -v --disable-pip-version-check --no-cache-dir ./
+```
+A Python-only build omits:
+- Fused kernels required to use `apex.optimizers.FusedAdam`.
+- Fused kernels required to use `apex.normalization.FusedLayerNorm` and `apex.normalization.FusedRMSNorm`.
+- Fused kernels that improve the performance and numerical stability of `apex.parallel.SyncBatchNorm`.
+- Fused kernels that improve the performance of `apex.parallel.DistributedDataParallel` and `apex.amp`.
+`DistributedDataParallel`, `amp`, and `SyncBatchNorm` will still be usable, but they may be slower.
+
+
+### [Experimental] Windows
+`pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" .` may work if you were able to build Pytorch from source
+on your system. A Python-only build via `pip install -v --no-cache-dir .` is more likely to work.
+If you installed Pytorch in a Conda environment, make sure to install Apex in that same environment.
+
+
+## Custom C++/CUDA Extensions and Install Options
+
+If a requirement of a module is not met, then it will not be built.
+
+| Module Name | Install Option | Misc |
+|---------------|------------------|--------|
+| `apex_C` | `--cpp_ext` | |
+| `amp_C` | `--cuda_ext` | |
+| `syncbn` | `--cuda_ext` | |
+| `fused_layer_norm_cuda` | `--cuda_ext` | [`apex.normalization`](./apex/normalization) |
+| `mlp_cuda` | `--cuda_ext` | |
+| `scaled_upper_triang_masked_softmax_cuda` | `--cuda_ext` | |
+| `generic_scaled_masked_softmax_cuda` | `--cuda_ext` | |
+| `scaled_masked_softmax_cuda` | `--cuda_ext` | |
+| `fused_weight_gradient_mlp_cuda` | `--cuda_ext` | Requires CUDA>=11 |
+| `permutation_search_cuda` | `--permutation_search` | [`apex.contrib.sparsity`](./apex/contrib/sparsity) |
+| `bnp` | `--bnp` | [`apex.contrib.groupbn`](./apex/contrib/groupbn) |
+| `xentropy` | `--xentropy` | [`apex.contrib.xentropy`](./apex/contrib/xentropy) |
+| `focal_loss_cuda` | `--focal_loss` | [`apex.contrib.focal_loss`](./apex/contrib/focal_loss) |
+| `fused_index_mul_2d` | `--index_mul_2d` | [`apex.contrib.index_mul_2d`](./apex/contrib/index_mul_2d) |
+| `fused_adam_cuda` | `--deprecated_fused_adam` | [`apex.contrib.optimizers`](./apex/contrib/optimizers) |
+| `fused_lamb_cuda` | `--deprecated_fused_lamb` | [`apex.contrib.optimizers`](./apex/contrib/optimizers) |
+| `fast_layer_norm` | `--fast_layer_norm` | [`apex.contrib.layer_norm`](./apex/contrib/layer_norm). different from `fused_layer_norm` |
+| `fmhalib` | `--fmha` | [`apex.contrib.fmha`](./apex/contrib/fmha) |
+| `fast_multihead_attn` | `--fast_multihead_attn` | [`apex.contrib.multihead_attn`](./apex/contrib/multihead_attn) |
+| `transducer_joint_cuda` | `--transducer` | [`apex.contrib.transducer`](./apex/contrib/transducer) |
+| `transducer_loss_cuda` | `--transducer` | [`apex.contrib.transducer`](./apex/contrib/transducer) |
+| `cudnn_gbn_lib` | `--cudnn_gbn` | Requires cuDNN>=8.5, [`apex.contrib.cudnn_gbn`](./apex/contrib/cudnn_gbn) |
+| `peer_memory_cuda` | `--peer_memory` | [`apex.contrib.peer_memory`](./apex/contrib/peer_memory) |
+| `nccl_p2p_cuda` | `--nccl_p2p` | Requires NCCL >= 2.10, [`apex.contrib.nccl_p2p`](./apex/contrib/nccl_p2p) |
+| `fast_bottleneck` | `--fast_bottleneck` | Requires `peer_memory_cuda` and `nccl_p2p_cuda`, [`apex.contrib.bottleneck`](./apex/contrib/bottleneck) |
+| `fused_conv_bias_relu` | `--fused_conv_bias_relu` | Requires cuDNN>=8.4, [`apex.contrib.conv_bias_relu`](./apex/contrib/conv_bias_relu) |
diff --git a/apex/apex/RNN/README.md b/apex/apex/RNN/README.md
new file mode 100644
index 00000000..82c4eb68
--- /dev/null
+++ b/apex/apex/RNN/README.md
@@ -0,0 +1,3 @@
+**This module will be removed by the end of February 2023**
+
+Under construction...
diff --git a/apex/apex/RNN/RNNBackend.py b/apex/apex/RNN/RNNBackend.py
new file mode 100644
index 00000000..a9382e60
--- /dev/null
+++ b/apex/apex/RNN/RNNBackend.py
@@ -0,0 +1,365 @@
+import torch
+import torch.nn as nn
+from torch.autograd import Variable
+
+import torch.nn.functional as F
+
+import math
+
+
+def is_iterable(maybe_iterable):
+ return isinstance(maybe_iterable, list) or isinstance(maybe_iterable, tuple)
+
+
+def flatten_list(tens_list):
+ """
+ flatten_list
+ """
+ if not is_iterable(tens_list):
+ return tens_list
+
+ return torch.cat(tens_list, dim=0).view(len(tens_list), *tens_list[0].size() )
+
+
+#These modules always assumes batch_first
+class bidirectionalRNN(nn.Module):
+ """
+ bidirectionalRNN
+ """
+ def __init__(self, inputRNN, num_layers=1, dropout = 0):
+ super(bidirectionalRNN, self).__init__()
+ self.dropout = dropout
+ self.fwd = stackedRNN(inputRNN, num_layers=num_layers, dropout = dropout)
+ self.bckwrd = stackedRNN(inputRNN.new_like(), num_layers=num_layers, dropout = dropout)
+ self.rnns = nn.ModuleList([self.fwd, self.bckwrd])
+
+ #collect hidden option will return all hidden/cell states from entire RNN
+ def forward(self, input, collect_hidden=False):
+ """
+ forward()
+ """
+ seq_len = input.size(0)
+ bsz = input.size(1)
+
+ fwd_out, fwd_hiddens = list(self.fwd(input, collect_hidden = collect_hidden))
+ bckwrd_out, bckwrd_hiddens = list(self.bckwrd(input, reverse=True, collect_hidden = collect_hidden))
+
+ output = torch.cat( [fwd_out, bckwrd_out], -1 )
+ hiddens = tuple( torch.cat(hidden, -1) for hidden in zip( fwd_hiddens, bckwrd_hiddens) )
+
+ return output, hiddens
+
+ def reset_parameters(self):
+ """
+ reset_parameters()
+ """
+ for rnn in self.rnns:
+ rnn.reset_parameters()
+
+ def init_hidden(self, bsz):
+ """
+ init_hidden()
+ """
+ for rnn in self.rnns:
+ rnn.init_hidden(bsz)
+
+ def detach_hidden(self):
+ """
+ detach_hidden()
+ """
+ for rnn in self.rnns:
+ rnn.detachHidden()
+
+ def reset_hidden(self, bsz):
+ """
+ reset_hidden()
+ """
+ for rnn in self.rnns:
+ rnn.reset_hidden(bsz)
+
+ def init_inference(self, bsz):
+ """
+ init_inference()
+ """
+ for rnn in self.rnns:
+ rnn.init_inference(bsz)
+
+
+#assumes hidden_state[0] of inputRNN is output hidden state
+#constructor either takes an RNNCell or list of RNN layers
+class stackedRNN(nn.Module):
+ """
+ stackedRNN
+ """
+ def __init__(self, inputRNN, num_layers=1, dropout=0):
+ super(stackedRNN, self).__init__()
+
+ self.dropout = dropout
+
+ if isinstance(inputRNN, RNNCell):
+ self.rnns = [inputRNN]
+ for i in range(num_layers-1):
+ self.rnns.append(inputRNN.new_like(inputRNN.output_size))
+ elif isinstance(inputRNN, list):
+ assert len(inputRNN) == num_layers, "RNN list length must be equal to num_layers"
+ self.rnns=inputRNN
+ else:
+ raise RuntimeError()
+
+ self.nLayers = len(self.rnns)
+
+ self.rnns = nn.ModuleList(self.rnns)
+
+
+ '''
+ Returns output as hidden_state[0] Tensor([sequence steps][batch size][features])
+ If collect hidden will also return Tuple(
+ [n_hidden_states][sequence steps] Tensor([layer][batch size][features])
+ )
+ If not collect hidden will also return Tuple(
+ [n_hidden_states] Tensor([layer][batch size][features])
+ '''
+ def forward(self, input, collect_hidden=False, reverse=False):
+ """
+ forward()
+ """
+ seq_len = input.size(0)
+ bsz = input.size(1)
+ inp_iter = reversed(range(seq_len)) if reverse else range(seq_len)
+
+ hidden_states = [[] for i in range(self.nLayers)]
+ outputs = []
+
+ for seq in inp_iter:
+ for layer in range(self.nLayers):
+
+ if layer == 0:
+ prev_out = input[seq]
+
+ outs = self.rnns[layer](prev_out)
+
+ if collect_hidden:
+ hidden_states[layer].append(outs)
+ elif seq == seq_len-1:
+ hidden_states[layer].append(outs)
+
+ prev_out = outs[0]
+
+ outputs.append(prev_out)
+
+ if reverse:
+ outputs = list(reversed(outputs))
+ '''
+ At this point outputs is in format:
+ list( [seq_length] x Tensor([bsz][features]) )
+ need to convert it to:
+ list( Tensor([seq_length][bsz][features]) )
+ '''
+ output = flatten_list(outputs)
+
+ '''
+ hidden_states at this point is in format:
+ list( [layer][seq_length][hidden_states] x Tensor([bsz][features]) )
+ need to convert it to:
+ For not collect hidden:
+ list( [hidden_states] x Tensor([layer][bsz][features]) )
+ For collect hidden:
+ list( [hidden_states][seq_length] x Tensor([layer][bsz][features]) )
+ '''
+ if not collect_hidden:
+ seq_len = 1
+ n_hid = self.rnns[0].n_hidden_states
+ new_hidden = [ [ [ None for k in range(self.nLayers)] for j in range(seq_len) ] for i in range(n_hid) ]
+
+
+ for i in range(n_hid):
+ for j in range(seq_len):
+ for k in range(self.nLayers):
+ new_hidden[i][j][k] = hidden_states[k][j][i]
+
+ hidden_states = new_hidden
+ #Now in format list( [hidden_states][seq_length][layer] x Tensor([bsz][features]) )
+ #Reverse seq_length if reverse
+ if reverse:
+ hidden_states = list( list(reversed(list(entry))) for entry in hidden_states)
+
+ #flatten layer dimension into tensor
+ hiddens = list( list(
+ flatten_list(seq) for seq in hidden )
+ for hidden in hidden_states )
+
+ #Now in format list( [hidden_states][seq_length] x Tensor([layer][bsz][features]) )
+ #Remove seq_length dimension if not collect_hidden
+ if not collect_hidden:
+ hidden_states = list( entry[0] for entry in hidden_states)
+ return output, hidden_states
+
+ def reset_parameters(self):
+ """
+ reset_parameters()
+ """
+ for rnn in self.rnns:
+ rnn.reset_parameters()
+
+ def init_hidden(self, bsz):
+ """
+ init_hidden()
+ """
+ for rnn in self.rnns:
+ rnn.init_hidden(bsz)
+
+ def detach_hidden(self):
+ """
+ detach_hidden()
+ """
+ for rnn in self.rnns:
+ rnn.detach_hidden()
+
+ def reset_hidden(self, bsz):
+ """
+ reset_hidden()
+ """
+ for rnn in self.rnns:
+ rnn.reset_hidden(bsz)
+
+ def init_inference(self, bsz):
+ """
+ init_inference()
+ """
+ for rnn in self.rnns:
+ rnn.init_inference(bsz)
+
+class RNNCell(nn.Module):
+ """
+ RNNCell
+ gate_multiplier is related to the architecture you're working with
+ For LSTM-like it will be 4 and GRU-like will be 3.
+ Always assumes input is NOT batch_first.
+ Output size that's not hidden size will use output projection
+ Hidden_states is number of hidden states that are needed for cell
+ if one will go directly to cell as tensor, if more will go as list
+ """
+ def __init__(self, gate_multiplier, input_size, hidden_size, cell, n_hidden_states = 2, bias = False, output_size = None):
+ super(RNNCell, self).__init__()
+
+ self.gate_multiplier = gate_multiplier
+ self.input_size = input_size
+ self.hidden_size = hidden_size
+ self.cell = cell
+ self.bias = bias
+ self.output_size = output_size
+ if output_size is None:
+ self.output_size = hidden_size
+
+ self.gate_size = gate_multiplier * self.hidden_size
+ self.n_hidden_states = n_hidden_states
+
+ self.w_ih = nn.Parameter(torch.empty(self.gate_size, self.input_size))
+ self.w_hh = nn.Parameter(torch.empty(self.gate_size, self.output_size))
+
+ #Check if there's recurrent projection
+ if(self.output_size != self.hidden_size):
+ self.w_ho = nn.Parameter(torch.empty(self.output_size, self.hidden_size))
+
+ self.b_ih = self.b_hh = None
+ if self.bias:
+ self.b_ih = nn.Parameter(torch.empty(self.gate_size))
+ self.b_hh = nn.Parameter(torch.empty(self.gate_size))
+
+ #hidden states for forward
+ self.hidden = [ None for states in range(self.n_hidden_states)]
+
+ self.reset_parameters()
+
+ def new_like(self, new_input_size=None):
+ """
+ new_like()
+ """
+ if new_input_size is None:
+ new_input_size = self.input_size
+
+ return type(self)(self.gate_multiplier,
+ new_input_size,
+ self.hidden_size,
+ self.cell,
+ self.n_hidden_states,
+ self.bias,
+ self.output_size)
+
+
+ #Use xavier where we can (weights), otherwise use uniform (bias)
+ def reset_parameters(self, gain=1):
+ """
+ reset_parameters()
+ """
+ stdev = 1.0 / math.sqrt(self.hidden_size)
+ for param in self.parameters():
+ param.data.uniform_(-stdev, stdev)
+ '''
+ Xavier reset:
+ def reset_parameters(self, gain=1):
+ stdv = 1.0 / math.sqrt(self.gate_size)
+
+ for param in self.parameters():
+ if (param.dim() > 1):
+ torch.nn.init.xavier_normal(param, gain)
+ else:
+ param.data.uniform_(-stdv, stdv)
+ '''
+ def init_hidden(self, bsz):
+ """
+ init_hidden()
+ """
+ for param in self.parameters():
+ if param is not None:
+ a_param = param
+ break
+
+ for i, _ in enumerate(self.hidden):
+ if(self.hidden[i] is None or self.hidden[i].data.size()[0] != bsz):
+
+ if i==0:
+ hidden_size = self.output_size
+ else:
+ hidden_size = self.hidden_size
+
+ tens = a_param.data.new(bsz, hidden_size).zero_()
+ self.hidden[i] = Variable(tens, requires_grad=False)
+
+
+ def reset_hidden(self, bsz):
+ """
+ reset_hidden()
+ """
+ for i, _ in enumerate(self.hidden):
+ self.hidden[i] = None
+ self.init_hidden(bsz)
+
+ def detach_hidden(self):
+ """
+ detach_hidden()
+ """
+ for i, _ in enumerate(self.hidden):
+ if self.hidden[i] is None:
+ raise RuntimeError("Must initialize hidden state before you can detach it")
+ for i, _ in enumerate(self.hidden):
+ self.hidden[i] = self.hidden[i].detach()
+
+ def forward(self, input):
+ """
+ forward()
+ if not inited or bsz has changed this will create hidden states
+ """
+ self.init_hidden(input.size()[0])
+
+ hidden_state = self.hidden[0] if self.n_hidden_states == 1 else self.hidden
+ self.hidden = self.cell(input, hidden_state, self.w_ih, self.w_hh, b_ih=self.b_ih, b_hh=self.b_hh)
+ if(self.n_hidden_states > 1):
+ self.hidden = list(self.hidden)
+ else:
+ self.hidden=[self.hidden]
+
+ if self.output_size != self.hidden_size:
+ self.hidden[0] = F.linear(self.hidden[0], self.w_ho)
+
+ return tuple(self.hidden)
diff --git a/apex/apex/RNN/__init__.py b/apex/apex/RNN/__init__.py
new file mode 100644
index 00000000..d7067466
--- /dev/null
+++ b/apex/apex/RNN/__init__.py
@@ -0,0 +1,3 @@
+from .models import LSTM, GRU, ReLU, Tanh, mLSTM
+
+__all__ = ['models']
diff --git a/apex/apex/RNN/cells.py b/apex/apex/RNN/cells.py
new file mode 100644
index 00000000..09b08581
--- /dev/null
+++ b/apex/apex/RNN/cells.py
@@ -0,0 +1,84 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .RNNBackend import RNNCell
+
+from torch.nn._functions.thnn import rnnFusedPointwise as fusedBackend
+
+import math
+
+
+class mLSTMRNNCell(RNNCell):
+ """
+ mLSTMRNNCell
+ """
+
+ def __init__(self, input_size, hidden_size, bias = False, output_size = None):
+ gate_multiplier = 4
+ super(mLSTMRNNCell, self).__init__(gate_multiplier, input_size, hidden_size, mLSTMCell, n_hidden_states = 2, bias = bias, output_size = output_size)
+
+ self.w_mih = nn.Parameter(torch.empty(self.output_size, self.input_size))
+ self.w_mhh = nn.Parameter(torch.empty(self.output_size, self.output_size))
+
+ self.reset_parameters()
+
+ def forward(self, input):
+ """
+ mLSTMRNNCell.forward()
+ """
+ #if not inited or bsz has changed this will create hidden states
+ self.init_hidden(input.size()[0])
+
+ hidden_state = self.hidden[0] if self.n_hidden_states == 1 else self.hidden
+
+ self.hidden = list(
+ self.cell(input, hidden_state, self.w_ih, self.w_hh, self.w_mih, self.w_mhh,
+ b_ih=self.b_ih, b_hh=self.b_hh)
+ )
+
+ if self.output_size != self.hidden_size:
+ self.hidden[0] = F.linear(self.hidden[0], self.w_ho)
+ return tuple(self.hidden)
+
+
+ def new_like(self, new_input_size=None):
+ if new_input_size is None:
+ new_input_size = self.input_size
+
+ return type(self)(
+ new_input_size,
+ self.hidden_size,
+ self.bias,
+ self.output_size)
+
+def mLSTMCell(input, hidden, w_ih, w_hh, w_mih, w_mhh, b_ih=None, b_hh=None):
+ """
+ mLSTMCell
+ """
+
+ if input.is_cuda:
+ igates = F.linear(input, w_ih)
+ m = F.linear(input, w_mih) * F.linear(hidden[0], w_mhh)
+ hgates = F.linear(m, w_hh)
+
+ state = fusedBackend.LSTMFused.apply
+ return state(igates, hgates, hidden[1], b_ih, b_hh)
+
+ hx, cx = hidden
+
+ m = F.linear(input, w_mih) * F.linear(hidden[0], w_mhh)
+ gates = F.linear(input, w_ih, b_ih) + F.linear(m, w_hh, b_hh)
+
+ ingate, forgetgate, cellgate, outgate = gates.chunk(4, 1)
+
+ ingate = F.sigmoid(ingate)
+ forgetgate = F.sigmoid(forgetgate)
+ cellgate = F.tanh(cellgate)
+ outgate = F.sigmoid(outgate)
+
+ cy = (forgetgate * cx) + (ingate * cellgate)
+ hy = outgate * F.tanh(cy)
+
+ return hy, cy
+
diff --git a/apex/apex/RNN/models.py b/apex/apex/RNN/models.py
new file mode 100644
index 00000000..d661aa0d
--- /dev/null
+++ b/apex/apex/RNN/models.py
@@ -0,0 +1,56 @@
+import torch
+
+from torch.nn._functions.rnn import LSTMCell, RNNReLUCell, RNNTanhCell, GRUCell
+
+from apex import deprecated_warning
+from .RNNBackend import bidirectionalRNN, stackedRNN, RNNCell
+from .cells import mLSTMRNNCell, mLSTMCell
+
+def toRNNBackend(inputRNN, num_layers, bidirectional=False, dropout = 0):
+ """
+ :class:`toRNNBackend`
+ """
+
+ deprecated_warning("`apex.RNN` is deprecated and will be removed by the end of February 2023.")
+ if bidirectional:
+ return bidirectionalRNN(inputRNN, num_layers, dropout = dropout)
+ else:
+ return stackedRNN(inputRNN, num_layers, dropout = dropout)
+
+
+def LSTM(input_size, hidden_size, num_layers, bias=True, batch_first=False, dropout=0, bidirectional=False, output_size = None):
+ """
+ :class:`LSTM`
+ """
+ inputRNN = RNNCell(4, input_size, hidden_size, LSTMCell, 2, bias, output_size)
+ return toRNNBackend(inputRNN, num_layers, bidirectional, dropout=dropout)
+
+def GRU(input_size, hidden_size, num_layers, bias=True, batch_first=False, dropout=0, bidirectional=False, output_size = None):
+ """
+ :class:`GRU`
+ """
+ inputRNN = RNNCell(3, input_size, hidden_size, GRUCell, 1, bias, output_size)
+ return toRNNBackend(inputRNN, num_layers, bidirectional, dropout=dropout)
+
+def ReLU(input_size, hidden_size, num_layers, bias=True, batch_first=False, dropout=0, bidirectional=False, output_size = None):
+ """
+ :class:`ReLU`
+ """
+ inputRNN = RNNCell(1, input_size, hidden_size, RNNReLUCell, 1, bias, output_size)
+ return toRNNBackend(inputRNN, num_layers, bidirectional, dropout=dropout)
+
+def Tanh(input_size, hidden_size, num_layers, bias=True, batch_first=False, dropout=0, bidirectional=False, output_size = None):
+ """
+ :class:`Tanh`
+ """
+ inputRNN = RNNCell(1, input_size, hidden_size, RNNTanhCell, 1, bias, output_size)
+ return toRNNBackend(inputRNN, num_layers, bidirectional, dropout=dropout)
+
+def mLSTM(input_size, hidden_size, num_layers, bias=True, batch_first=False, dropout=0, bidirectional=False, output_size = None):
+ """
+ :class:`mLSTM`
+ """
+ inputRNN = mLSTMRNNCell(input_size, hidden_size, bias=bias, output_size=output_size)
+ return toRNNBackend(inputRNN, num_layers, bidirectional, dropout=dropout)
+
+
diff --git a/apex/apex/__init__.py b/apex/apex/__init__.py
new file mode 100644
index 00000000..74851f5b
--- /dev/null
+++ b/apex/apex/__init__.py
@@ -0,0 +1,68 @@
+import logging
+import warnings
+
+# May help avoid undefined symbol errors https://pytorch.org/cppdocs/notes/faq.html#undefined-symbol-errors-from-pytorch-aten
+import torch
+
+
+__all__ = ["amp", "fp16_utils", "optimizers", "normalization", "transformer"]
+
+
+if torch.distributed.is_available():
+ from . import parallel
+ __all__.append("parallel")
+
+from . import amp
+from . import fp16_utils
+
+# For optimizers and normalization there is no Python fallback.
+# Absence of cuda backend is a hard error.
+# I would like the errors from importing fused_adam_cuda or fused_layer_norm_cuda
+# to be triggered lazily, because if someone has installed with --cpp_ext and --cuda_ext
+# so they expect those backends to be available, but for some reason they actually aren't
+# available (for example because they built improperly in a way that isn't revealed until
+# load time) the error message is timely and visible.
+from . import optimizers
+from . import normalization
+from . import transformer
+
+
+# Logging utilities for apex.transformer module
+class RankInfoFormatter(logging.Formatter):
+
+ def format(self, record):
+ from apex.transformer.parallel_state import get_rank_info
+ record.rank_info = get_rank_info()
+ return super().format(record)
+
+
+_library_root_logger = logging.getLogger(__name__)
+handler = logging.StreamHandler()
+handler.setFormatter(RankInfoFormatter("%(asctime)s - PID:%(process)d - rank:%(rank_info)s - %(filename)s:%(lineno)d - %(levelname)s - %(message)s", "%y-%m-%d %H:%M:%S"))
+_library_root_logger.addHandler(handler)
+_library_root_logger.propagate = False
+
+
+def check_cudnn_version_and_warn(global_option: str, required_cudnn_version: int) -> bool:
+ cudnn_available = torch.backends.cudnn.is_available()
+ cudnn_version = torch.backends.cudnn.version() if cudnn_available else None
+ if not (cudnn_available and (cudnn_version >= required_cudnn_version)):
+ warnings.warn(
+ f"`{global_option}` depends on cuDNN {required_cudnn_version} or later, "
+ f"but {'cuDNN is not available' if not cudnn_available else cudnn_version}"
+ )
+ return False
+ return True
+
+
+class DeprecatedFeatureWarning(FutureWarning):
+ pass
+
+
+def deprecated_warning(msg: str) -> None:
+ if (
+ not torch.distributed.is_available
+ or not torch.distributed.is_initialized()
+ or (torch.distributed.is_initialized() and torch.distributed.get_rank() == 0)
+ ):
+ warnings.warn(msg, DeprecatedFeatureWarning)
diff --git a/apex/apex/_autocast_utils.py b/apex/apex/_autocast_utils.py
new file mode 100644
index 00000000..3a92a83f
--- /dev/null
+++ b/apex/apex/_autocast_utils.py
@@ -0,0 +1,26 @@
+from typing import Optional, Sequence
+
+import torch
+
+
+__all__ = ["_cast_if_autocast_enabled"]
+
+
+def _get_autocast_dtypes() -> Sequence[torch.dtype]:
+ if torch.cuda.is_bf16_supported():
+ return [torch.half, torch.bfloat16]
+ return [torch.half]
+
+
+def _get_current_dtype(dtype: Optional[torch.dtype] = None) -> torch.dtype:
+ if not torch.is_autocast_enabled():
+ return torch.float or dtype
+ else:
+ return torch.get_autocast_gpu_dtype()
+
+
+def _cast_if_autocast_enabled(*args):
+ if not torch.is_autocast_enabled():
+ return args
+ else:
+ return torch.cuda.amp.autocast_mode._cast(args, torch.get_autocast_gpu_dtype())
diff --git a/apex/apex/amp/README.md b/apex/apex/amp/README.md
new file mode 100644
index 00000000..a87b5010
--- /dev/null
+++ b/apex/apex/amp/README.md
@@ -0,0 +1,72 @@
+# amp: Automatic Mixed Precision
+
+## Annotating User Functions
+
+Nearly all PyTorch user code needs nothing more than the two steps
+above to use amp. After all, custom layers are built out of simpler
+PyTorch components, and amp already can see those.
+
+However, any custom C++ or CUDA code is outside of amp's (default)
+view of things. For example, suppose I implemented a new recurrent
+cell called a "forgetful recurrent unit" that calls directly into a
+CUDA backend:
+
+```python
+from backend import FRUBackend
+
+def fru(input, hidden, weight, bias):
+ # call to CUDA code
+ FRUBackend(input, hidden, weight, bias)
+```
+
+In this case, it is possible to get a runtime type mismatch. For
+example, you might have `input` in fp16, and `weight` in fp32, and amp
+doesn't have the visibility to insert an appropriate cast.
+
+amp exposes two ways to handle "invisible" backend code: function
+annotations and explicit registration.
+
+#### Function annotation
+
+The first way to handle backend code is a set of function annotations:
+
+- `@amp.half_function`
+- `@amp.float_function`
+- `@amp.promote_function`
+
+These correspond to:
+
+- Cast all arguments to fp16
+- Cast all argumnets fo fp32
+- If there are any type mismatches, cast everything to the widest type
+
+In our example, we believe that the FRU unit is fp16-safe and will get
+performance gains from casting its arguments to fp16, so we write:
+
+```python
+@amp.half_function
+def fru(input, hidden, weight, bias):
+ #...
+```
+
+#### Explicit registration
+
+The other way to handle backend code is with explicit function
+registration:
+
+- `amp.register_half_function(module, function_name)`
+- `amp.register_float_function(module, function_name)`
+- `amp.register_promote_function(module, function_name)`
+
+When using this API, `module` is the containing class or module for
+the function, and `function_name` is the _string_ name of the
+function. Note that the function must be registered before the call to
+`amp.initalize()`.
+
+For our FRU unit, we can register the backend function directly:
+
+```python
+import backend
+
+amp.register_half_function(backend, 'FRUBackend')
+```
diff --git a/apex/apex/amp/__init__.py b/apex/apex/amp/__init__.py
new file mode 100644
index 00000000..34d080a6
--- /dev/null
+++ b/apex/apex/amp/__init__.py
@@ -0,0 +1,5 @@
+from .amp import init, half_function, float_function, promote_function,\
+ register_half_function, register_float_function, register_promote_function
+from .handle import scale_loss, disable_casts
+from .frontend import initialize, state_dict, load_state_dict
+from ._amp_state import master_params, _amp_state
diff --git a/apex/apex/amp/__version__.py b/apex/apex/amp/__version__.py
new file mode 100644
index 00000000..3a83701b
--- /dev/null
+++ b/apex/apex/amp/__version__.py
@@ -0,0 +1,2 @@
+VERSION = (0, 1, 0)
+__version__ = '.'.join(map(str, VERSION))
diff --git a/apex/apex/amp/_amp_state.py b/apex/apex/amp/_amp_state.py
new file mode 100644
index 00000000..7e8a329f
--- /dev/null
+++ b/apex/apex/amp/_amp_state.py
@@ -0,0 +1,59 @@
+# This is a "header object" that allows different amp modules to communicate.
+# I'm a C++ guy, not a python guy. I decided this approach because it seemed most C++-like.
+# But apparently it's ok:
+# http://effbot.org/pyfaq/how-do-i-share-global-variables-across-modules.htm
+import torch
+
+
+class AmpState(object):
+ def __init__(self):
+ self.hard_override=False
+ self.allow_incoming_model_not_fp32 = False
+ self.verbosity=1
+
+
+# Attribute stash. Could also just stash things as global module attributes.
+_amp_state = AmpState()
+
+
+def warn_or_err(msg):
+ if _amp_state.hard_override:
+ print("Warning: " + msg)
+ else:
+ raise RuntimeError(msg)
+ # I'm not sure if allowing hard_override is a good idea.
+ # + " If you're sure you know what you're doing, supply " +
+ # "hard_override=True to amp.initialize.")
+
+
+def maybe_print(msg, rank0=False):
+ distributed = torch.distributed.is_available() and \
+ torch.distributed.is_initialized() and \
+ torch.distributed.get_world_size() > 1
+ if _amp_state.verbosity > 0:
+ if rank0:
+ if distributed:
+ if torch.distributed.get_rank() == 0:
+ print(msg)
+ else:
+ print(msg)
+ else:
+ print(msg)
+
+
+# def iter_params(param_groups):
+# for group in param_groups:
+# for p in group['params']:
+# yield p
+
+
+def master_params(optimizer):
+ """
+ Generator expression that iterates over the params owned by ``optimizer``.
+
+ Args:
+ optimizer: An optimizer previously returned from ``amp.initialize``.
+ """
+ for group in optimizer.param_groups:
+ for p in group['params']:
+ yield p
diff --git a/apex/apex/amp/_initialize.py b/apex/apex/amp/_initialize.py
new file mode 100644
index 00000000..3ae6fded
--- /dev/null
+++ b/apex/apex/amp/_initialize.py
@@ -0,0 +1,265 @@
+import collections.abc as container_abcs
+from types import MethodType
+import functools
+import sys
+import warnings
+
+import numpy as np
+import torch
+
+from ._amp_state import _amp_state, warn_or_err
+from .handle import disable_casts
+from .scaler import LossScaler
+from ._process_optimizer import _process_optimizer
+from apex.fp16_utils import convert_network
+from ..fp16_utils import FP16_Optimizer as FP16_Optimizer_general
+from ..contrib.optimizers import FP16_Optimizer as FP16_Optimizer_for_fused
+
+if torch.distributed.is_available():
+ from ..parallel import DistributedDataParallel as apex_DDP
+ from ..parallel.LARC import LARC
+
+
+def to_type(dtype, t):
+ if isinstance(t, torch.Tensor):
+ if not t.is_cuda:
+ # This should not be a hard error, since it may be legitimate.
+ warnings.warn("An input tensor was not cuda.")
+ # GANs require this.
+ # if t.requires_grad:
+ # warn_or_err("input data requires grad. Since input data is not a model parameter,\n"
+ # "its gradients will not be properly allreduced by DDP.")
+ if t.is_floating_point():
+ return t.to(dtype)
+ return t
+ else:
+ # Trust the user's custom batch type, that's all I can do here.
+ return t.to(dtype)
+
+
+# Modified from torch.optim.optimizer.py. This is a bit more general than casted_args in utils.py.
+def applier(value, fn):
+ if isinstance(value, torch.Tensor):
+ return fn(value)
+ elif isinstance(value, str):
+ return value
+ elif isinstance(value, np.ndarray):
+ return value
+ elif hasattr(value, "to"): # Allow handling of custom batch classes
+ return fn(value)
+ elif isinstance(value, container_abcs.Mapping):
+ return {applier(k, fn) : applier(v, fn) for k, v in value.items()}
+ elif isinstance(value, container_abcs.Iterable):
+ return type(value)(applier(v, fn) for v in value)
+ else:
+ # Do I want this to fire off even if someone chooses to pass something ordinary like
+ # an int or float? May be more annoying than it's worth.
+ # print("Warning: unrecognized type in applier. If your input data is a custom class, "
+ # "provide it with a .to(dtype) method which converts its floating-point Tensors to dtype. "
+ # "Amp will check for your custom to() and invoke it to cast the batch's "
+ # "floating-point Tensors to the appropriate type. "
+ # "Also, if your data is a custom class, it is your responsibility to ensure that "
+ # "any Tensors you want to be cuda are already cuda."
+ return value
+
+
+def check_models(models):
+ for model in models:
+ parallel_type = None
+ if isinstance(model, torch.nn.parallel.DistributedDataParallel):
+ parallel_type = "torch.nn.parallel.DistributedDataParallel"
+ if ('apex_DDP' in sys.modules) and isinstance(model, apex_DDP):
+ parallel_type = "apex.parallel.DistributedDataParallel"
+ if isinstance(model, torch.nn.parallel.DataParallel):
+ parallel_type = "torch.nn.parallel.DataParallel"
+ if parallel_type is not None:
+ raise RuntimeError("Incoming model is an instance of {}. ".format(parallel_type) +
+ "Parallel wrappers should only be applied to the model(s) AFTER \n"
+ "the model(s) have been returned from amp.initialize.")
+
+
+def check_params_fp32(models):
+ for model in models:
+ for name, param in model.named_parameters():
+ if param.is_floating_point():
+ if 'Half' in param.type():
+ warn_or_err("Found param {} with type {}, expected torch.cuda.FloatTensor.\n"
+ "When using amp.initialize, you do not need to call .half() on your model\n"
+ "before passing it, no matter what optimization level you choose.".format(
+ name, param.type()))
+ elif not param.is_cuda:
+ warn_or_err("Found param {} with type {}, expected torch.cuda.FloatTensor.\n"
+ "When using amp.initialize, you need to provide a model with parameters\n"
+ "located on a CUDA device before passing it no matter what optimization level\n"
+ "you chose. Use model.to('cuda') to use the default device.".format(
+ name, param.type()))
+
+ # Backward compatibility for PyTorch 0.4
+ if hasattr(model, 'named_buffers'):
+ buf_iter = model.named_buffers()
+ else:
+ buf_iter = model._buffers
+ for obj in buf_iter:
+ if type(obj)==tuple:
+ name, buf = obj
+ else:
+ name, buf = obj, buf_iter[obj]
+ if buf.is_floating_point():
+ if 'Half' in buf.type():
+ warn_or_err("Found buffer {} with type {}, expected torch.cuda.FloatTensor.\n"
+ "When using amp.initialize, you do not need to call .half() on your model\n"
+ "before passing it, no matter what optimization level you choose.".format(
+ name, buf.type()))
+ elif not buf.is_cuda:
+ warn_or_err("Found buffer {} with type {}, expected torch.cuda.FloatTensor.\n"
+ "When using amp.initialize, you need to provide a model with buffers\n"
+ "located on a CUDA device before passing it no matter what optimization level\n"
+ "you chose. Use model.to('cuda') to use the default device.".format(
+ name, buf.type()))
+
+
+def check_optimizers(optimizers):
+ for optim in optimizers:
+ bad_optim_type = None
+ if isinstance(optim, FP16_Optimizer_general):
+ bad_optim_type = "apex.fp16_utils.FP16_Optimizer"
+ if isinstance(optim, FP16_Optimizer_for_fused):
+ bad_optim_type = "apex.optimizers.FP16_Optimizer"
+ if bad_optim_type is not None:
+ raise RuntimeError("An incoming optimizer is an instance of {}. ".format(bad_optim_type) +
+ "The optimizer(s) passed to amp.initialize() must be bare \n"
+ "instances of either ordinary Pytorch optimizers, or Apex fused \n"
+ "optimizers.\n")
+
+
+class O2StateDictHook(object):
+ def __init__(self, fn):
+ self.fn = fn
+
+ def __call__(self, module, state_dict, prefix, local_metadata):
+ for key in state_dict:
+ param = state_dict[key]
+ if 'Half' in param.type():
+ param = param.to(torch.float32)
+ state_dict[key] = param
+
+
+def _initialize(models, optimizers, properties, num_losses=1, cast_model_outputs=None):
+ from .amp import init as amp_init
+
+ optimizers_was_list = False
+ if isinstance(optimizers, torch.optim.Optimizer) or ('LARC' in globals() and isinstance(optimizers, LARC)):
+ optimizers = [optimizers]
+ elif optimizers is None:
+ optimizers = []
+ elif isinstance(optimizers, list):
+ optimizers_was_list = True
+ check_optimizers(optimizers)
+ else:
+ check_optimizers([optimizers])
+ raise TypeError("optimizers must be either a single optimizer or a list of optimizers.")
+
+ if isinstance(models, torch.nn.Module):
+ models_was_list = False
+ models = [models]
+ elif isinstance(models, list):
+ models_was_list = True
+ else:
+ raise TypeError("models must be either a single model or a list of models.")
+
+ check_models(models)
+
+ if not _amp_state.allow_incoming_model_not_fp32:
+ check_params_fp32(models)
+
+ # In the future, when FP16_Optimizer can be deprecated and master weights can
+ # become an attribute, remember to stash master weights before casting the model.
+
+ if properties.cast_model_type:
+ if properties.keep_batchnorm_fp32:
+ for model in models:
+ convert_network(model, properties.cast_model_type)
+ else:
+ for model in models:
+ model.to(properties.cast_model_type)
+
+ input_caster = functools.partial(to_type, properties.cast_model_type)
+ if cast_model_outputs is not None:
+ output_caster = functools.partial(to_type, cast_model_outputs)
+ else:
+ output_caster = functools.partial(to_type, torch.float32)
+
+ for model in models:
+ # Patch the forward method to cast incoming data to the correct type, and
+ # outgoing data to float32, so "the user never needs to call .half()."
+ # I like writing things explicitly more than decorators.
+ def patch_forward(old_fwd):
+ def new_fwd(*args, **kwargs):
+ output = old_fwd(*applier(args, input_caster),
+ **applier(kwargs, input_caster))
+ return applier(output, output_caster)
+ return new_fwd
+
+ model.forward = patch_forward(model.forward)
+
+ # State dict trick to recast any preexisting per-param state tensors
+ for optimizer in optimizers:
+ optimizer.load_state_dict(optimizer.state_dict())
+
+ # patch model.state_dict() to return float32 params
+ for model in models:
+ for module in model.modules():
+ module._register_state_dict_hook(O2StateDictHook(functools.partial(to_type, torch.float32)))
+
+ elif cast_model_outputs is not None:
+ output_caster = functools.partial(to_type, cast_model_outputs)
+
+ for model in models:
+ def patch_forward(old_fwd):
+ def new_fwd(*args, **kwargs):
+ output = old_fwd(*args, **kwargs)
+ return applier(output, output_caster)
+ return new_fwd
+
+ model.forward = patch_forward(model.forward)
+
+ for i, optimizer in enumerate(optimizers):
+ optimizers[i] = _process_optimizer(optimizer, properties)
+
+ _amp_state.loss_scalers = []
+ for _ in range(num_losses):
+ _amp_state.loss_scalers.append(LossScaler(properties.loss_scale,
+ min_loss_scale=_amp_state.min_loss_scale,
+ max_loss_scale=_amp_state.max_loss_scale))
+
+ if properties.patch_torch_functions:
+ # handle is unused here. It's accessible later through a global value anyway.
+ handle = amp_init(loss_scale=properties.loss_scale, verbose=(_amp_state.verbosity == 2))
+ for optimizer in optimizers:
+ # Disable Amp casting for the optimizer step, because it should only be
+ # applied to FP32 master params anyway.
+ def patch_step(old_step):
+ def new_step(self, *args, **kwargs):
+ with disable_casts():
+ output = old_step(*args, **kwargs)
+ return output
+ return new_step
+
+ optimizer.step = MethodType(patch_step(optimizer.step), optimizer)
+
+ if optimizers_was_list:
+ if models_was_list:
+ return models, optimizers
+ else:
+ return models[0], optimizers
+ else:
+ if models_was_list:
+ if len(optimizers) == 0:
+ return models
+ else:
+ return models, optimizers[0]
+ else:
+ if len(optimizers) == 0:
+ return models[0]
+ else:
+ return models[0], optimizers[0]
diff --git a/apex/apex/amp/_process_optimizer.py b/apex/apex/amp/_process_optimizer.py
new file mode 100644
index 00000000..471289bb
--- /dev/null
+++ b/apex/apex/amp/_process_optimizer.py
@@ -0,0 +1,489 @@
+import types
+from ..fp16_utils import master_params_to_model_params
+from ..multi_tensor_apply import multi_tensor_applier
+from ._amp_state import maybe_print
+import torch
+from ..optimizers import FusedSGD
+
+
+class AmpOptimizerState(object):
+ def __init__(self):
+ pass
+
+
+def _master_params_to_model_params(self):
+ stash = self._amp_stash
+ if multi_tensor_applier.available:
+ if len(stash.all_fp16_params) > 0:
+ multi_tensor_applier(
+ stash.multi_tensor_scale,
+ stash.dummy_overflow_buf,
+ [stash.all_fp32_from_fp16_params, stash.all_fp16_params],
+ 1.0)
+ else:
+ for fp16_group, fp32_from_fp16_group in zip(stash.fp16_groups, stash.fp32_from_fp16_groups):
+ master_params_to_model_params(fp16_group, fp32_from_fp16_group)
+
+
+def lazy_init_with_master_weights(self):
+ stash = self._amp_stash
+ stash.fp16_groups = []
+ stash.fp32_from_fp16_groups = []
+ stash.fp32_from_fp32_groups = []
+ for i, param_group in enumerate(self.param_groups):
+ # maybe_print("FP16_Optimizer processing param group {}:".format(i))
+ fp16_params_this_group = []
+ fp32_params_this_group = []
+ fp32_from_fp16_params_this_group = []
+ for i, param in enumerate(param_group['params']):
+ if param.requires_grad:
+ if param.type() == 'torch.cuda.HalfTensor':
+ # maybe_print("FP16_Optimizer received torch.cuda.HalfTensor with {}"
+ # .format(param.size()))
+ fp16_params_this_group.append(param)
+ master_param = param.detach().clone().float()
+ master_param.requires_grad = True
+ param_group['params'][i] = master_param
+ fp32_from_fp16_params_this_group.append(master_param)
+ # Reset existing state dict key to the new master param.
+ # We still need to recast per-param state tensors, if any, to FP32.
+ if param in self.state:
+ self.state[master_param] = self.state.pop(param)
+ elif param.type() == 'torch.cuda.FloatTensor':
+ # maybe_print("FP16_Optimizer received torch.cuda.FloatTensor with {}"
+ # .format(param.size()))
+ fp32_params_this_group.append(param)
+ param_group['params'][i] = param
+ else:
+ raise TypeError("Optimizer's parameters must be either "
+ "torch.cuda.FloatTensor or torch.cuda.HalfTensor. "
+ "Received {}".format(param.type()))
+
+ stash.fp16_groups.append(fp16_params_this_group)
+ stash.fp32_from_fp16_groups.append(fp32_from_fp16_params_this_group)
+ stash.fp32_from_fp32_groups.append(fp32_params_this_group)
+
+ stash.all_fp16_params = []
+ for group in stash.fp16_groups:
+ stash.all_fp16_params += group
+
+ stash.all_fp32_from_fp16_params = []
+ for group in stash.fp32_from_fp16_groups:
+ stash.all_fp32_from_fp16_params += group
+
+ stash.all_fp32_from_fp32_params = []
+ for group in stash.fp32_from_fp32_groups:
+ stash.all_fp32_from_fp32_params += group
+
+ # all_fp16_grad_stash is only needed for fused optimizers.
+ stash.all_fp16_grad_stash = [None for _ in stash.all_fp16_params]
+ # stash.all_fp32_from_fp16_grad_stash = [None for _ in stash.all_fp32_from_fp16_params]
+ stash.all_fp32_from_fp32_grad_stash = [None for _ in stash.all_fp32_from_fp32_params]
+
+ for param in stash.all_fp32_from_fp16_params:
+ param.grad = None
+
+ for param in stash.all_fp32_from_fp32_params:
+ param.grad = None
+
+ # Leverage state_dict() and load_state_dict() to recast preexisting per-param state tensors
+ self.load_state_dict(self.state_dict())
+
+
+def post_backward_models_are_masters(scaler, params, stashed_grads, scale_override=None):
+ grads_have_scale, stashed_have_scale, out_scale = scaler.loss_scale(), 1.0, 1.0
+
+ # not much to do if scale == 1.0 and static scaling
+ if scaler.loss_scale() == 1.0 and not scaler.dynamic:
+ # Clear the stash.
+ for i in range(len(stashed_grads)):
+ stashed_grads[i] = None
+ return
+
+ if scale_override is not None:
+ grads_have_scale, stashed_have_scale, out_scale = scale_override
+
+ # This is a lot of python overhead...
+ grads_needing_unscale = []
+ grads_needing_unscale_with_stash = []
+ stashed = []
+ for param, stashed_grad in zip(params, stashed_grads):
+ if param.grad is None and stashed_grad is not None:
+ param.grad = stashed_grad
+ elif param.grad is not None and stashed_grad is None:
+ grads_needing_unscale.append(param.grad)
+ elif param.grad is not None and stashed_grad is not None:
+ grads_needing_unscale_with_stash.append(param.grad)
+ stashed.append(stashed_grad)
+ else: # param.grad is None and stashed_grad is None
+ continue
+
+ # unscale() implements grads*(1/scale), so "scale" should be grads_have_scale/out_scale.
+ if len(grads_needing_unscale) > 0:
+ scaler.unscale(
+ grads_needing_unscale,
+ grads_needing_unscale,
+ None, # unused_scale, currently present to avoid API breakage elsewhere
+ models_are_masters=True,
+ scale_override=grads_have_scale/out_scale)
+
+ if len(grads_needing_unscale_with_stash) > 0:
+ scaler.unscale_with_stashed(
+ grads_needing_unscale_with_stash,
+ stashed,
+ grads_needing_unscale_with_stash,
+ scale_override=(grads_have_scale, stashed_have_scale, out_scale))
+
+ # Clear the stash.
+ for i in range(len(stashed_grads)):
+ stashed_grads[i] = None
+
+
+def prepare_backward_with_master_weights(self):
+ stash = self._amp_stash
+
+ self._amp_lazy_init()
+
+ for i, param in enumerate(stash.all_fp16_params):
+ # Set up to leverage grad copy elision.
+ # This may behave differently from an unpatched optimizer if zero_grad is used and the param is unused.
+ param.grad = None
+
+ # for i, param in enumerate(stash.all_fp32_from_fp16_params):
+ # stash.all_fp32_from_fp16_grad_stash[i] = param.grad
+
+ for i, param in enumerate(stash.all_fp32_from_fp32_params):
+ stash.all_fp32_from_fp32_grad_stash[i] = param.grad
+ # Set up to leverage grad copy elision:
+ param.grad = None
+
+
+def post_backward_with_master_weights(self, scaler):
+ stash = self._amp_stash
+
+ self._amp_lazy_init()
+
+ # This is a lot of python overhead...
+ fp16_grads_needing_unscale = []
+ new_fp32_grads = []
+ fp16_grads_needing_unscale_with_stash = []
+ preexisting_fp32_grads = []
+ for fp16_param, fp32_param in zip(stash.all_fp16_params,
+ stash.all_fp32_from_fp16_params):
+ if fp16_param.grad is None and fp32_param.grad is not None:
+ continue
+ elif fp16_param.grad is not None and fp32_param.grad is None:
+ fp32_param.grad = torch.empty_like(fp32_param)
+ fp16_grads_needing_unscale.append(fp16_param.grad)
+ new_fp32_grads.append(fp32_param.grad)
+ elif fp16_param.grad is not None and fp32_param.grad is not None:
+ fp16_grads_needing_unscale_with_stash.append(fp16_param.grad)
+ preexisting_fp32_grads.append(fp32_param.grad)
+ else: # fp16_param.grad is None and fp32_param.grad is None:
+ continue
+
+ if len(fp16_grads_needing_unscale) > 0:
+ scaler.unscale(
+ fp16_grads_needing_unscale,
+ new_fp32_grads,
+ scaler.loss_scale(),
+ models_are_masters=False)
+
+ if len(fp16_grads_needing_unscale_with_stash) > 0:
+ scaler.unscale_with_stashed(
+ fp16_grads_needing_unscale_with_stash,
+ preexisting_fp32_grads,
+ preexisting_fp32_grads)
+
+ # fp32 params can be treated as they would be in the "no_master_weights" case.
+ post_backward_models_are_masters(
+ scaler,
+ stash.all_fp32_from_fp32_params,
+ stash.all_fp32_from_fp32_grad_stash)
+
+
+def lazy_init_no_master_weights(self):
+ stash = self._amp_stash
+ stash.all_fp16_params = []
+ stash.all_fp32_params = []
+ for i, param_group in enumerate(self.param_groups):
+ for i, param in enumerate(param_group['params']):
+ if param.type() == 'torch.cuda.HalfTensor':
+ stash.all_fp16_params.append(param)
+ elif param.type() == 'torch.cuda.FloatTensor':
+ stash.all_fp32_params.append(param)
+ else:
+ raise TypeError("Optimizer's parameters must be either "
+ "torch.cuda.FloatTensor or torch.cuda.HalfTensor. "
+ "Received {}".format(param.type()))
+
+ stash.all_fp16_grad_stash = [None for _ in stash.all_fp16_params]
+ stash.all_fp32_grad_stash = [None for _ in stash.all_fp32_params]
+
+
+def prepare_backward_no_master_weights(self):
+ stash = self._amp_stash
+
+ self._amp_lazy_init()
+
+ for i, param in enumerate(stash.all_fp16_params):
+ stash.all_fp16_grad_stash[i] = param.grad
+ # Set up to leverage grad copy elision:
+ param.grad = None
+
+ for i, param in enumerate(stash.all_fp32_params):
+ stash.all_fp32_grad_stash[i] = param.grad
+ # Set up to leverage grad copy elision:
+ param.grad = None
+
+
+def post_backward_no_master_weights(self, scaler):
+ stash = self._amp_stash
+
+ self._amp_lazy_init()
+
+ split_types = ((stash.all_fp16_params, stash.all_fp16_grad_stash),
+ (stash.all_fp32_params, stash.all_fp32_grad_stash))
+
+ for params, stashed_grads in split_types:
+ post_backward_models_are_masters(scaler, params, stashed_grads)
+
+
+#####################################################################################
+# FusedSGD versions
+#####################################################################################
+
+# FusedSGD never explicitly materializes the fp32 gradients for "fp32 from fp16" master params
+# outside the kernel, so we must accumulate directly into the model grads.
+def prepare_backward_with_master_weights_FusedSGD(self):
+ if self.materialize_master_grads:
+ prepare_backward_with_master_weights(self)
+ else:
+ stash = self._amp_stash
+
+ self._amp_lazy_init()
+
+ for i, param in enumerate(stash.all_fp16_params):
+ stash.all_fp16_grad_stash[i] = param.grad
+ # Set up to leverage grad copy elision:
+ param.grad = None
+
+ for i, param in enumerate(stash.all_fp32_from_fp32_params):
+ stash.all_fp32_from_fp32_grad_stash[i] = param.grad
+ # Set up to leverage grad copy elision:
+ param.grad = None
+
+
+def post_backward_with_master_weights_FusedSGD(self, scaler):
+ if self.materialize_master_grads:
+ post_backward_with_master_weights(self, scaler)
+ else:
+ stash = self._amp_stash
+
+ self._amp_lazy_init()
+
+ grads_have_scale = scaler.loss_scale()
+ stashed_have_scale = self.most_recent_scale
+ out_scale = grads_have_scale
+ if self.scale_set_by_backward:
+ out_scale = min(grads_have_scale, self.most_recent_scale)
+
+ split_types = ((stash.all_fp16_params, stash.all_fp16_grad_stash),
+ (stash.all_fp32_from_fp32_params, stash.all_fp32_from_fp32_grad_stash))
+
+
+ # unscale_with_stashed() implements grads*1/scale + stashed_grads*1.
+ # stashed_grads are scaled by self.most_recent_scale.
+ for params, stashed_grads in split_types:
+ post_backward_models_are_masters(scaler, params, stashed_grads,
+ (grads_have_scale, stashed_have_scale, out_scale))
+
+ self.most_recent_scale = out_scale
+ self.scale_set_by_backward = True
+
+
+def prepare_backward_no_master_weights_FusedSGD(self):
+ prepare_backward_no_master_weights(self)
+
+
+def post_backward_no_master_weights_FusedSGD(self, scaler):
+ post_backward_no_master_weights(self, scaler)
+
+
+def _amp_lazy_init(self):
+ stash = self._amp_stash
+
+ if not stash.lazy_init_called:
+ self._lazy_init_maybe_master_weights()
+ stash.lazy_init_called = True
+
+
+def _process_optimizer(optimizer, properties):
+ if hasattr(optimizer, "_amp_stash"):
+ raise RuntimeError("A given optimizer should only be passed through amp.initialize once.")
+ else:
+ optimizer._amp_stash = AmpOptimizerState()
+
+ optimizer._amp_stash.lazy_init_called = False
+ optimizer._amp_stash.already_patched = False
+ optimizer._amp_stash.params_have_scaled_gradients = False
+
+ for name in ("_lazy_init_maybe_master_weights",
+ "_master_params_to_model_params",
+ "_prepare_amp_backward",
+ "_post_amp_backward",
+ "_amp_lazy_init"):
+ if hasattr(optimizer, name):
+ raise RuntimeError("Incoming optimizer already has {} defined.".format(name))
+
+ # TODO: Centralize exposure and import error checking for the C backend.
+ if multi_tensor_applier.available:
+ import amp_C
+ optimizer._amp_stash.multi_tensor_scale = amp_C.multi_tensor_scale
+ optimizer._amp_stash.multi_tensor_l2norm = amp_C.multi_tensor_l2norm
+ optimizer._amp_stash.dummy_overflow_buf = torch.cuda.IntTensor([0]);
+
+ if properties.master_weights:
+ optimizer._lazy_init_maybe_master_weights = types.MethodType(
+ lazy_init_with_master_weights, optimizer)
+
+ optimizer._master_params_to_model_params = types.MethodType(
+ _master_params_to_model_params, optimizer)
+
+ old_step = optimizer.step
+ def new_step(self, closure=None):
+ if closure is not None:
+ raise RuntimeError("Currently, Amp does not support closure use with optimizers.")
+ retval = old_step()
+ if not isinstance(self, FusedSGD):
+ self._master_params_to_model_params()
+ # Clear the master grads that wouldn't be zeroed by model.zero_grad()
+ for param in self._amp_stash.all_fp32_from_fp16_params:
+ param.grad = None
+ return retval
+ optimizer.step = types.MethodType(new_step, optimizer)
+
+ old_zero_grad = optimizer.zero_grad
+ def new_zero_grad(self):
+ stash = self._amp_stash
+ self._amp_lazy_init()
+ # Zero the model grads.
+ for param in stash.all_fp16_params:
+ if param.grad is not None:
+ param.grad.detach_()
+ param.grad.zero_()
+ for param in stash.all_fp32_from_fp32_params:
+ if param.grad is not None:
+ param.grad.detach_()
+ param.grad.zero_()
+ # Clear the master grads that are independent of model grads
+ for param in self._amp_stash.all_fp32_from_fp16_params:
+ param.grad = None
+ optimizer.zero_grad = types.MethodType(new_zero_grad, optimizer)
+
+ if isinstance(optimizer, FusedSGD):
+ optimizer._prepare_amp_backward = types.MethodType(
+ prepare_backward_with_master_weights_FusedSGD, optimizer)
+ optimizer._post_amp_backward = types.MethodType(
+ post_backward_with_master_weights_FusedSGD, optimizer)
+ else:
+ optimizer._prepare_amp_backward = types.MethodType(
+ prepare_backward_with_master_weights, optimizer)
+ optimizer._post_amp_backward = types.MethodType(
+ post_backward_with_master_weights, optimizer)
+ else:
+ optimizer._lazy_init_maybe_master_weights = types.MethodType(
+ lazy_init_no_master_weights, optimizer)
+
+ if isinstance(optimizer, FusedSGD):
+ optimizer._prepare_amp_backward = types.MethodType(
+ prepare_backward_no_master_weights_FusedSGD, optimizer)
+ optimizer._post_amp_backward = types.MethodType(
+ post_backward_no_master_weights_FusedSGD, optimizer)
+ else:
+ optimizer._prepare_amp_backward = types.MethodType(
+ prepare_backward_no_master_weights, optimizer)
+ optimizer._post_amp_backward = types.MethodType(
+ post_backward_no_master_weights, optimizer)
+
+ optimizer._amp_lazy_init = types.MethodType(_amp_lazy_init, optimizer)
+
+ old_add_param_group = optimizer.add_param_group
+
+ def new_add_param_group(self, new_group):
+ stash = self._amp_stash
+
+ if not stash.lazy_init_called:
+ self._lazy_init_maybe_master_weights()
+ stash.lazy_init_called = True
+
+ assert isinstance(new_group, dict), "param group must be a dict"
+
+ new_params = new_group['params']
+ if isinstance(new_params, torch.Tensor):
+ new_group['params'] = [new_params]
+ elif isinstance(new_params, set):
+ raise TypeError('optimizer parameters need to be organized in ordered collections, but '
+ 'the ordering of tensors in sets will change between runs. Please use a list instead.')
+ else:
+ new_group['params'] = list(new_params)
+
+ if properties.master_weights:
+ # Mutate new_group in-place to use FP32 master params
+ fp16_params_this_group = []
+ fp32_params_this_group = []
+ fp32_from_fp16_params_this_group = []
+ for i, param in enumerate(new_group['params']):
+ if param.requires_grad:
+ if param.type() == 'torch.cuda.HalfTensor':
+ fp16_params_this_group.append(param)
+ master_param = param.detach().clone().float()
+ master_param.requires_grad = True
+ new_group['params'][i] = master_param
+ fp32_from_fp16_params_this_group.append(master_param)
+ elif param.type() == 'torch.cuda.FloatTensor':
+ fp32_params_this_group.append(param)
+ new_group['params'][i] = param
+ else:
+ raise TypeError("Optimizer's parameters must be either "
+ "torch.cuda.FloatTensor or torch.cuda.HalfTensor. "
+ "Received {}".format(param.type()))
+
+ stash.fp16_groups.append(fp16_params_this_group)
+ stash.fp32_from_fp16_groups.append(fp32_from_fp16_params_this_group)
+ stash.fp32_from_fp32_groups.append(fp32_params_this_group)
+
+ stash.all_fp16_params += fp16_params_this_group
+ stash.all_fp32_from_fp16_params += fp32_from_fp16_params_this_group
+ stash.all_fp32_from_fp32_params += fp32_params_this_group
+
+ # stash.all_fp32_from_fp16_grad_stash = [None for _ in stash.all_fp32_from_fp16_params]
+ stash.all_fp32_from_fp32_grad_stash += [None for _ in fp32_params_this_group]
+
+ # It should be ok to let params be added with existing .grad attributes.
+ # for param in fp16_params_this_group:
+ # param.grad = None
+
+ # for param in fp32_from_fp16_params_this_group:
+ # param.grad = None
+
+ # for param in stash.fp32_params_this_group:
+ # param.grad = None
+ else:
+ for param in new_group['params']:
+ if param.type() == 'torch.cuda.HalfTensor':
+ stash.all_fp16_params.append(param)
+ stash.all_fp16_grad_stash.append(None)
+ elif param.type() == 'torch.cuda.FloatTensor':
+ stash.all_fp32_params.append(param)
+ stash.all_fp32_grad_stash.append(None)
+ else:
+ raise TypeError("Optimizer's parameters must be either "
+ "torch.cuda.FloatTensor or torch.cuda.HalfTensor. "
+ "Received {}".format(param.type()))
+
+ old_add_param_group(new_group)
+
+ optimizer.add_param_group = types.MethodType(new_add_param_group, optimizer)
+
+ return optimizer
diff --git a/apex/apex/amp/amp.py b/apex/apex/amp/amp.py
new file mode 100644
index 00000000..1a604666
--- /dev/null
+++ b/apex/apex/amp/amp.py
@@ -0,0 +1,183 @@
+import functools
+import itertools
+
+import torch
+
+from . import compat, rnn_compat, utils, wrap
+from .handle import AmpHandle, NoOpHandle
+from .lists import functional_overrides, torch_overrides, tensor_overrides
+from ._amp_state import _amp_state
+from .frontend import *
+
+
+_DECORATOR_HANDLE = None
+_USER_CAST_REGISTRY = set()
+_USER_PROMOTE_REGISTRY = set()
+
+
+def _decorator_helper(orig_fn, cast_fn, wrap_fn):
+ def wrapper(*args, **kwargs):
+ handle = _DECORATOR_HANDLE
+ if handle is None or not handle.is_active():
+ return orig_fn(*args, **kwargs)
+ inner_cast_fn = utils.verbosify(cast_fn, orig_fn.__name__,
+ handle.verbose)
+ return wrap_fn(orig_fn, inner_cast_fn, handle)(*args, **kwargs)
+ return wrapper
+
+
+# Decorator form
+def half_function(fn):
+ from apex import deprecated_warning
+ deprecated_warning("apex.amp is deprecated and will be removed by the end of February 2023. Use [PyTorch AMP](https://pytorch.org/docs/stable/amp.html)")
+ wrap_fn = functools.partial(wrap.make_cast_wrapper, try_caching=True)
+ return _decorator_helper(fn, utils.maybe_half, wrap_fn)
+
+
+def float_function(fn):
+ from apex import deprecated_warning
+ deprecated_warning("apex.amp is deprecated and will be removed by the end of February 2023. Use [PyTorch AMP](https://pytorch.org/docs/stable/amp.html)")
+ wrap_fn = functools.partial(wrap.make_cast_wrapper, try_caching=False)
+ return _decorator_helper(fn, utils.maybe_float, wrap_fn)
+
+
+def promote_function(fn):
+ from apex import deprecated_warning
+ deprecated_warning("apex.amp is deprecated and will be removed by the end of February 2023. Use [PyTorch AMP](https://pytorch.org/docs/stable/amp.html)")
+ wrap_fn = functools.partial(wrap.make_promote_wrapper)
+ return _decorator_helper(fn, utils.maybe_float, wrap_fn)
+
+
+# Registry form
+def register_half_function(module, name):
+ if not hasattr(module, name):
+ raise ValueError('No function named {} in module {}.'.format(
+ name, module))
+ _USER_CAST_REGISTRY.add((module, name, utils.maybe_half))
+
+
+def register_float_function(module, name):
+ if not hasattr(module, name):
+ raise ValueError('No function named {} in module {}.'.format(
+ name, module))
+ _USER_CAST_REGISTRY.add((module, name, utils.maybe_float))
+
+
+def register_promote_function(module, name):
+ if not hasattr(module, name):
+ raise ValueError('No function named {} in module {}.'.format(
+ name, module))
+ _USER_PROMOTE_REGISTRY.add((module, name))
+
+
+# Top-level function to insert _all_ the hooks.
+def init(enabled=True, loss_scale="dynamic", enable_caching=True, verbose=False, allow_banned=False):
+ global _DECORATOR_HANDLE
+
+ if not enabled:
+ handle = NoOpHandle()
+ _DECORATOR_HANDLE = handle
+ return handle
+
+ handle = AmpHandle(loss_scale, enable_caching, verbose)
+
+ # 0) Force-{fp16, fp32} for user-annotated functions
+ for mod, fn, cast_fn in _USER_CAST_REGISTRY:
+ try_caching = (cast_fn == utils.maybe_half)
+ wrap.cached_cast(mod, fn, cast_fn, handle,
+ try_caching, verbose)
+ _USER_CAST_REGISTRY.clear()
+
+ # 0.5) Force-promote for user-annotated functions
+ for mod, fn in _USER_PROMOTE_REGISTRY:
+ wrap.promote(mod, fn, handle, verbose)
+ _USER_PROMOTE_REGISTRY.clear()
+
+ # 1) Force-{fp16, fp32} on white- / black-list functions
+ override_modules = [functional_overrides,
+ torch_overrides,
+ tensor_overrides]
+ cast_table = [('FP16_FUNCS', utils.maybe_half),
+ ('FP32_FUNCS', utils.maybe_float)]
+ for module, (list_name, cast_fn) in itertools.product(override_modules,
+ cast_table):
+ for fn in getattr(module, list_name):
+ try_caching = (cast_fn == utils.maybe_half)
+ wrap.cached_cast(module.MODULE, fn, cast_fn, handle,
+ try_caching, verbose)
+
+ # 1.5) Pre-0.4, put the blacklist methods on HalfTensor and whitelist
+ # methods on FloatTensor, since they're distinct types.
+ if compat.tensor_is_float_tensor():
+ for fn in tensor_overrides.FP16_FUNCS:
+ wrap.cached_cast(torch.cuda.FloatTensor, fn, utils.maybe_half,
+ handle, try_caching=True, verbose=verbose)
+ for fn in tensor_overrides.FP32_FUNCS:
+ wrap.cached_cast(torch.cuda.HalfTensor, fn, utils.maybe_float,
+ handle, try_caching=False, verbose=verbose)
+
+ # 2) Enable type-promotion on multi-arg functions and methods.
+ # NB: special handling for sequence fns (e.g. `torch.cat`).
+ promote_modules = [torch_overrides, tensor_overrides]
+ promote_table = [('CASTS', wrap.promote),
+ ('SEQUENCE_CASTS', wrap.sequence_promote)]
+ for promote_mod, (list_name, promote_fn) in itertools.product(promote_modules,
+ promote_table):
+ for fn in getattr(promote_mod, list_name):
+ promote_fn(promote_mod.MODULE, fn, handle, verbose)
+
+ # 2.5) Pre-0.4, add blacklist methods directly to HalfTensor and FloatTensor types
+ if compat.tensor_is_float_tensor():
+ for cls, (list_name, promote_fn) in itertools.product([torch.cuda.FloatTensor,
+ torch.cuda.HalfTensor],
+ promote_table):
+ for fn in getattr(tensor_overrides, list_name):
+ promote_fn(cls, fn, handle, verbose)
+
+ # 3) For any in-place version of a blacklist function, error if any input is fp16.
+ # NB: this is overly conservative.
+ for fn in utils.as_inplace(torch_overrides.FP32_FUNCS):
+ wrap.err_if_any_half(torch_overrides.MODULE, fn, handle)
+
+ # 3.5) For any in-place blacklist method, error if called on fp16 tensor
+ for fn in utils.as_inplace(tensor_overrides.FP32_FUNCS):
+ wrap.err_if_arg0_half(tensor_overrides.MODULE, fn, handle, verbose)
+ if compat.tensor_is_float_tensor():
+ wrap.err_if_arg0_half(torch.cuda.HalfTensor, fn, handle, verbose)
+
+ # 4) For other in-place methods, match the type of self tensor
+ for fn in utils.as_inplace(itertools.chain(
+ tensor_overrides.FP16_FUNCS,
+ tensor_overrides.CASTS)):
+ wrap.promote_match_arg0(tensor_overrides.MODULE, fn, handle, verbose)
+ if compat.tensor_is_float_tensor():
+ wrap.promote_match_arg0(torch.cuda.HalfTensor, fn, handle, verbose)
+ wrap.promote_match_arg0(torch.cuda.FloatTensor, fn, handle, verbose)
+
+ # 5) RNNs + RNN cells are whitelisted specially
+ if rnn_compat.has_old_rnns():
+ wrap.rnn_cast(torch.nn.backends.thnn.backend, 'RNN', handle, verbose)
+ if not rnn_compat.has_old_rnns():
+ # Patch in our own indirection of `_VF` in modules/rnn s.t. it is mutable.
+ torch.nn.modules.rnn._VF = rnn_compat.VariableFunctionsShim()
+ # Wrap all the rnns
+ for x in rnn_compat.RNN_NAMES:
+ wrap.new_rnn_cast(x.upper(), handle, verbose)
+
+ # Wrap all the RNN cells
+ rnn_compat.whitelist_rnn_cells(handle, verbose)
+
+ # 6) Place error+print message on banned functions.
+ # Or, if allow_banned, then cast to FP32.
+ for fn, err_msg in functional_overrides.BANNED_FUNCS:
+ if allow_banned:
+ wrap.cached_cast(functional_overrides.MODULE, fn, utils.maybe_float,
+ handle, try_caching=True, verbose=verbose)
+ else:
+ wrap.err_if_any_half(functional_overrides.MODULE, fn, handle, err_msg)
+
+ _DECORATOR_HANDLE = handle
+
+ _amp_state.handle = handle
+
+ return handle
diff --git a/apex/apex/amp/compat.py b/apex/apex/amp/compat.py
new file mode 100644
index 00000000..22276bd4
--- /dev/null
+++ b/apex/apex/amp/compat.py
@@ -0,0 +1,46 @@
+import torch
+
+# True for post-0.4, when Variables/Tensors merged.
+def variable_is_tensor():
+ v = torch.autograd.Variable()
+ return isinstance(v, torch.Tensor)
+
+def tensor_is_variable():
+ x = torch.Tensor()
+ return type(x) == torch.autograd.Variable
+
+# False for post-0.4
+def tensor_is_float_tensor():
+ x = torch.Tensor()
+ return type(x) == torch.FloatTensor
+
+# Akin to `torch.is_tensor`, but returns True for Variable
+# objects in pre-0.4.
+def is_tensor_like(x):
+ return torch.is_tensor(x) or isinstance(x, torch.autograd.Variable)
+
+# Wraps `torch.is_floating_point` if present, otherwise checks
+# the suffix of `x.type()`.
+def is_floating_point(x):
+ if hasattr(torch, 'is_floating_point'):
+ return torch.is_floating_point(x)
+ try:
+ torch_type = x.type()
+ return torch_type.endswith('FloatTensor') or \
+ torch_type.endswith('HalfTensor') or \
+ torch_type.endswith('DoubleTensor')
+ except AttributeError:
+ return False
+
+def scalar_python_val(x):
+ if hasattr(x, 'item'):
+ return x.item()
+ else:
+ if isinstance(x, torch.autograd.Variable):
+ return x.data[0]
+ else:
+ return x[0]
+
+# Accounts for the possibility that some ops may be removed from a namespace.
+def filter_attrs(module, attrs):
+ return list(attrname for attrname in attrs if hasattr(module, attrname))
diff --git a/apex/apex/amp/frontend.py b/apex/apex/amp/frontend.py
new file mode 100644
index 00000000..616fb113
--- /dev/null
+++ b/apex/apex/amp/frontend.py
@@ -0,0 +1,446 @@
+from collections import OrderedDict
+
+import torch
+
+from ._initialize import _initialize
+from ._amp_state import _amp_state, warn_or_err, maybe_print
+
+
+class Properties(object):
+ """
+ This class has two purposes: to establish a set of default properties,
+ and to route setting of these attributes through __setattr__ so that (in theory)
+ they can be checked for consistency with other existing args.
+ """
+ def __init__(self):
+ self.options = {
+ "enabled" : False,
+ "opt_level" : None,
+ "cast_model_type" : None,
+ "patch_torch_functions" : False,
+ "keep_batchnorm_fp32" : None,
+ "master_weights" : None,
+ "loss_scale" : 1.0,
+ # Reserved for future functionality
+ # "fused_optimizer" : False,
+ # "enable_ddp_interop" : False,
+ }
+
+ """
+ This function allows updating several options at a time without routing through
+ __setattr__ checks, to avoid "you can't get there from here" scenarios.
+ Currently not intended to be exposed; users are expected to select an opt_level
+ and apply consistent modifications.
+ """
+ def _update_options_dict(self, new_options):
+ for k, v in new_options:
+ if k in self.options:
+ self.options[k] = v
+ else:
+ raise ValueError("Tried to set unexpected option {}".format(k))
+ """
+ The members of "options" are not direct attributes of self, so access attempts
+ will roll down to __getattr__. This borrows from the logic in torch.nn.Module.
+ """
+ def __getattr__(self, name):
+ if "options" in self.__dict__:
+ options = self.__dict__["options"]
+ if name in options:
+ return options[name]
+ raise AttributeError("'{}' object has no attribute '{}'".format(
+ type(self).__name__, name))
+
+ def __setattr__(self, name, value):
+ if "options" in self.__dict__:
+ if name in self.options:
+ # print("setting {} {}".format(name, value))
+ if name == "cast_model_type":
+ if self.opt_level == "O1" and value is not None:
+ if value is not False:
+ if value is not torch.float32:
+ warn_or_err("O1 inserts casts around Torch functions rather than "
+ "model weights, so with O1, the model weights themselves "
+ "should remain FP32. If you wish to cast the model to a "
+ "different type, use opt_level='O2' or 'O3'. " +
+ "cast_model_type was {}".format(value))
+ self.options[name] = value
+ elif name == "patch_torch_functions":
+ if self.opt_level != "O1" and value:
+ warn_or_err("Currently, patch_torch_functions=True should only be set by "
+ "selecting opt_level='O1'.")
+ self.options[name] = value
+ elif name == "keep_batchnorm_fp32":
+ if self.opt_level == "O1" and value is not None:
+ warn_or_err("With opt_level O1, batchnorm functions are automatically patched "
+ "to run in FP32, so keep_batchnorm_fp32 should be None." +
+ " keep_batchnorm_fp32 was {}".format(value))
+ if value == "False":
+ self.options[name] = False
+ elif value == "True":
+ self.options[name] = True
+ else:
+ assert (value is True or value is False or value is None),\
+ "keep_batchnorm_fp32 must be a boolean, the string 'True' or 'False', "\
+ "or None, found keep_batchnorm_fp32={}".format(value)
+ self.options[name] = value
+ elif name == "master_weights":
+ if self.opt_level == "O1" and value is not None:
+ warn_or_err("It doesn't make sense to use master_weights with O1. "
+ "With O1, your model weights themselves should be FP32.")
+ self.options[name] = value
+ elif name == "loss_scale":
+ if value == "dynamic":
+ self.options[name] = value
+ else:
+ self.options[name] = float(value)
+ else:
+ self.options[name] = value
+ else:
+ super(Properties, self).__setattr__(name, value)
+
+
+""" O0-O3 are convenience wrappers to establish defaults for typically used mixed precision options. """
+
+class O3:
+ brief = "O3: Pure FP16 training."
+ more = "Calls .half() on your model, converting the entire model to FP16.\n"\
+ "A casting operation is also inserted to cast incoming Tensors to FP16,\n"\
+ "so you don't need to change your data pipeline.\n"\
+ "This mode is useful for establishing a performance ceiling.\n"\
+ "It's also possible training may 'just work' in this mode.\n"\
+ "If not, try other optimization levels."
+
+ def __call__(self, properties):
+ properties.enabled = True
+ properties.opt_level = "O3"
+ properties.cast_model_type = torch.float16
+ properties.patch_torch_functions = False
+ properties.keep_batchnorm_fp32 = False
+ properties.master_weights = False
+ properties.loss_scale = 1.0
+ # properties.fused_optimizer = False
+ # properties.enable_ddp_interop = False
+ return properties # modified in place so this isn't really necessary
+
+
+class O2:
+ brief = "O2: FP16 training with FP32 batchnorm and FP32 master weights.\n"
+ more = "Calls .half() on your model, converting the entire model (except for batchnorms)\n"\
+ "to FP16. Batchnorms are retained in FP32 for additional stability.\n"\
+ "The forward pass is patched to cast incoming Tensors to FP16, so you don't need to change\n"\
+ "your data pipeline.\n"\
+ "O2 creates FP32 master weights outside the model and patches any optimizers to update\n"\
+ "these master weights, then copy the master weights into the FP16 model weights.\n"\
+ "Master weights can also improve convergence and stability."
+
+ def __call__(self, properties):
+ properties.enabled = True
+ properties.opt_level = "O2"
+ properties.cast_model_type = torch.float16
+ properties.patch_torch_functions = False
+ properties.keep_batchnorm_fp32 = True
+ properties.master_weights = True
+ properties.loss_scale = "dynamic"
+ # properties.fused_optimizer = False
+ # properties.enable_ddp_interop = False
+ return properties # modified in place so this isn't really necessary
+
+
+class O1:
+ brief = "O1: Insert automatic casts around Pytorch functions and Tensor methods.\n"
+ more = "The type of your model's weights is not altered. However, internally,\n"\
+ "Pytorch functions are patched to cast any Tensor Core-friendly ops to FP16 for speed,\n"\
+ "while operations that might benefit from the additional stability of FP32 are patched\n"\
+ "to cast their inputs to fp32.\n"\
+ "O1 is the safest way to try mixed precision training, and is recommended when\n"\
+ "trying mixed precision training for the first time."
+
+ def __call__(self, properties):
+ properties.enabled = True
+ properties.opt_level = "O1"
+ properties.cast_model_type = None
+ properties.patch_torch_functions = True
+ properties.keep_batchnorm_fp32 = None
+ properties.master_weights = None
+ properties.loss_scale = "dynamic"
+ # properties.fused_optimizer = False
+ # properties.enable_ddp_interop = False
+ return properties # modified in place so this isn't really necessary
+
+
+class O0:
+ brief = "O0: Pure FP32 training.\n"
+ more = "Your models are checked to make sure parameters are FP32, but otherwise the\n"\
+ "types of weights and internal Pytorch operations are not altered. This mode disables any\n"\
+ "FP16 arithmetic, although other optimizations like DDP interop may still be requested.\n"
+
+ def __call__(self, properties):
+ properties.enabled = True
+ properties.opt_level = "O0"
+ properties.cast_model_type = torch.float32
+ properties.patch_torch_functions = False
+ properties.keep_batchnorm_fp32 = None
+ properties.master_weights = False
+ properties.loss_scale = 1.0
+ # properties.fused_optimizer = False
+ # properties.enable_ddp_interop = False
+ return properties # modified in place so this isn't really necessary
+
+
+opt_levels = {"O3": O3(),
+ "O2": O2(),
+ "O1": O1(),
+ "O0": O0()}
+
+
+# allow user to directly pass Properties struct as well?
+def initialize(
+ models,
+ optimizers=None,
+ enabled=True,
+ opt_level="O1",
+ cast_model_type=None,
+ patch_torch_functions=None,
+ keep_batchnorm_fp32=None,
+ master_weights=None,
+ loss_scale=None,
+ cast_model_outputs=None,
+ num_losses=1,
+ verbosity=1,
+ min_loss_scale=None,
+ max_loss_scale=2.**24
+ ):
+ """
+ Initialize your models, optimizers, and the Torch tensor and functional namespace according to the
+ chosen ``opt_level`` and overridden properties, if any.
+
+ ``amp.initialize`` should be called **after** you have finished
+ constructing your model(s) and
+ optimizer(s), but **before** you send your model through any DistributedDataParallel wrapper.
+ See `Distributed training`_ in the Imagenet example.
+
+ Currently, ``amp.initialize`` should only be called **once**,
+ although it can process an arbitrary number of
+ models and optimizers (see the corresponding `Advanced Amp Usage topic`_).
+ If you think your use case requires ``amp.initialize`` to be called more than once,
+ `let us know`_.
+
+ Any property keyword argument that is not ``None`` will be interpreted as a manual override.
+
+ To prevent having to rewrite anything else in your script, name the returned models/optimizers
+ to replace the passed models/optimizers, as in the code sample below.
+
+ Args:
+ models (torch.nn.Module or list of torch.nn.Modules): Models to modify/cast.
+ optimizers (optional, torch.optim.Optimizer or list of torch.optim.Optimizers): Optimizers to modify/cast.
+ REQUIRED for training, optional for inference.
+ enabled (bool, optional, default=True): If False, renders all Amp calls no-ops, so your script
+ should run as if Amp were not present.
+ opt_level (str, optional, default="O1"): Pure or mixed precision optimization level. Accepted values are
+ "O0", "O1", "O2", and "O3", explained in detail above.
+ cast_model_type (``torch.dtype``, optional, default=None): Optional property override, see
+ above.
+ patch_torch_functions (bool, optional, default=None): Optional property override.
+ keep_batchnorm_fp32 (bool or str, optional, default=None): Optional property override. If
+ passed as a string, must be the string "True" or "False".
+ master_weights (bool, optional, default=None): Optional property override.
+ loss_scale (float or str, optional, default=None): Optional property override. If passed as a string,
+ must be a string representing a number, e.g., "128.0", or the string "dynamic".
+ cast_model_outputs (torch.dtype, optional, default=None): Option to ensure that the outputs
+ of your model(s) are always cast to a particular type regardless of ``opt_level``.
+ num_losses (int, optional, default=1): Option to tell Amp in advance how many losses/backward
+ passes you plan to use. When used in conjunction with the ``loss_id`` argument to
+ ``amp.scale_loss``, enables Amp to use a different loss scale per loss/backward pass,
+ which can improve stability. See "Multiple models/optimizers/losses"
+ under `Advanced Amp Usage`_ for examples. If ``num_losses`` is left to 1, Amp will still
+ support multiple losses/backward passes, but use a single global loss scale
+ for all of them.
+ verbosity (int, default=1): Set to 0 to suppress Amp-related output.
+ min_loss_scale (float, default=None): Sets a floor for the loss scale values that can be chosen by dynamic
+ loss scaling. The default value of None means that no floor is imposed.
+ If dynamic loss scaling is not used, `min_loss_scale` is ignored.
+ max_loss_scale (float, default=2.**24): Sets a ceiling for the loss scale values that can be chosen by
+ dynamic loss scaling. If dynamic loss scaling is not used, `max_loss_scale` is ignored.
+
+ Returns:
+ Model(s) and optimizer(s) modified according to the ``opt_level``.
+ If either the ``models`` or ``optimizers`` args were lists, the corresponding return value will
+ also be a list.
+
+ Permissible invocations::
+
+ model, optim = amp.initialize(model, optim,...)
+ model, [optim1, optim2] = amp.initialize(model, [optim1, optim2],...)
+ [model1, model2], optim = amp.initialize([model1, model2], optim,...)
+ [model1, model2], [optim1, optim2] = amp.initialize([model1, model2], [optim1, optim2],...)
+
+ # This is not an exhaustive list of the cross product of options that are possible,
+ # just a set of examples.
+ model, optim = amp.initialize(model, optim, opt_level="O0")
+ model, optim = amp.initialize(model, optim, opt_level="O0", loss_scale="dynamic"|128.0|"128.0")
+
+ model, optim = amp.initialize(model, optim, opt_level="O1") # uses "loss_scale="dynamic" default
+ model, optim = amp.initialize(model, optim, opt_level="O1", loss_scale=128.0|"128.0")
+
+ model, optim = amp.initialize(model, optim, opt_level="O2") # uses "loss_scale="dynamic" default
+ model, optim = amp.initialize(model, optim, opt_level="O2", loss_scale=128.0|"128.0")
+ model, optim = amp.initialize(model, optim, opt_level="O2", keep_batchnorm_fp32=True|False|"True"|"False")
+
+ model, optim = amp.initialize(model, optim, opt_level="O3") # uses loss_scale=1.0 default
+ model, optim = amp.initialize(model, optim, opt_level="O3", loss_scale="dynamic"|128.0|"128.0")
+ model, optim = amp.initialize(model, optim, opt_level="O3", keep_batchnorm_fp32=True|False|"True"|"False")
+
+ The `Imagenet example`_ demonstrates live use of various opt_levels and overrides.
+
+ .. _`Distributed training`:
+ https://github.com/NVIDIA/apex/tree/master/examples/imagenet#distributed-training
+
+ .. _`Imagenet example`:
+ https://github.com/NVIDIA/apex/tree/master/examples/imagenet
+
+ .. _`Advanced Amp Usage`:
+ https://nvidia.github.io/apex/advanced.html
+
+ .. _`Advanced Amp Usage topic`:
+ https://nvidia.github.io/apex/advanced.html#multiple-models-optimizers-losses
+
+ .. _`let us know`:
+ https://github.com/NVIDIA/apex/issues
+ """
+ from apex import deprecated_warning
+ deprecated_warning("apex.amp is deprecated and will be removed by the end of February 2023. Use [PyTorch AMP](https://pytorch.org/docs/stable/amp.html)")
+ _amp_state.opt_properties = Properties()
+ _amp_state.verbosity = verbosity
+
+ if not enabled:
+ if optimizers is None:
+ return models
+ else:
+ return models, optimizers
+
+ if not torch.backends.cudnn.enabled:
+ raise RuntimeError(
+ "Amp requires torch.backends.cudnn.enabled = True")
+
+ if opt_level not in opt_levels:
+ raise RuntimeError(
+ "Unexpected optimization level {}. ".format(opt_level) +
+ "Options are 'O0', 'O1', 'O2', 'O3'. Note that in `O0`, `O1`, etc., the prefix O is the letter O, " +
+ "not the number zero.")
+ else:
+ _amp_state.opt_properties = opt_levels[opt_level](_amp_state.opt_properties)
+ maybe_print("Selected optimization level {}".format(opt_levels[opt_level].brief), True)
+ maybe_print("Defaults for this optimization level are:", True)
+ for k, v in _amp_state.opt_properties.options.items():
+ maybe_print("{:22} : {}".format(k, v), True)
+
+ _amp_state.min_loss_scale = min_loss_scale
+ _amp_state.max_loss_scale = max_loss_scale
+
+ maybe_print("Processing user overrides (additional kwargs that are not None)...", True)
+ # I chose to have the keyword arguments listed directly in the argument list,
+ # instead of **kwargs, so I can't use kwargs.items() here.
+ if enabled is not None:
+ _amp_state.opt_properties.enabled = enabled
+ if opt_level is not None:
+ _amp_state.opt_properties.opt_level = opt_level
+ if cast_model_type is not None:
+ _amp_state.opt_properties.cast_model_type = cast_model_type
+ if patch_torch_functions is not None:
+ _amp_state.opt_properties.patch_torch_functions = patch_torch_functions
+ if keep_batchnorm_fp32 is not None:
+ _amp_state.opt_properties.keep_batchnorm_fp32 = keep_batchnorm_fp32
+ if master_weights is not None:
+ _amp_state.opt_properties.master_weights = master_weights
+ if loss_scale is not None:
+ _amp_state.opt_properties.loss_scale = loss_scale
+
+ maybe_print("After processing overrides, optimization options are:", True)
+ for k, v in _amp_state.opt_properties.options.items():
+ maybe_print("{:22} : {}".format(k, v), True)
+
+ return _initialize(models, optimizers, _amp_state.opt_properties, num_losses, cast_model_outputs)
+
+
+def state_dict(destination=None):
+ if destination is None:
+ destination = OrderedDict()
+
+ for idx, loss_scaler in enumerate(_amp_state.loss_scalers):
+ destination['loss_scaler%d' % idx] = {
+ 'loss_scale': loss_scaler.loss_scale(),
+ 'unskipped': loss_scaler._unskipped,
+ }
+ return destination
+
+
+def load_state_dict(state_dict):
+ # Check if state_dict containes the same number of loss_scalers as current setup
+ if len(state_dict) != len(_amp_state.loss_scalers):
+ print('Warning: state_dict contains {} entries, while {} loss_scalers are used'.format(
+ len(state_dict), len(_amp_state.loss_scalers)))
+
+ state_dict = state_dict.copy()
+
+ nb_loss_scalers = len(_amp_state.loss_scalers)
+ unexpected_keys = []
+ # Initialize idx outside, since unexpected_keys will increase it if enumerate is used
+ idx = 0
+ for key in state_dict:
+ if 'loss_scaler' not in key:
+ unexpected_keys.append(key)
+ else:
+ if idx > (nb_loss_scalers - 1):
+ print('Skipping loss_scaler[{}], since num_losses was set to {}'.format(
+ idx, nb_loss_scalers))
+ break
+ _amp_state.loss_scalers[idx]._loss_scale = state_dict[key]['loss_scale']
+ _amp_state.loss_scalers[idx]._unskipped = state_dict[key]['unskipped']
+ idx += 1
+
+ if len(unexpected_keys) > 0:
+ raise RuntimeError(
+ 'Error(s) in loading state_dict. Unexpected key(s) in state_dict: {}. '.format(
+ ', '.join('"{}"'.format(k) for k in unexpected_keys)))
+
+
+# TODO: is this necessary/useful?
+# def check_option_consistency(enabled=True,
+# opt_level=None,
+# cast_model_type=None,
+# patch_torch_functions=None,
+# keep_batchnorm_fp32=None,
+# master_weights=None,
+# loss_scale=None,
+# enable_ddp_interop=None,
+# hard_override=False):
+# """
+# Utility function that enables users to quickly check if the option combination they intend
+# to use is permitted. ``check_option_consistency`` does not require models or optimizers
+# to be constructed, and can be called at any point in the script. ``check_option_consistency``
+# is totally self-contained; it does not set any amp global state or affect anything outside
+# of itself.
+# """
+#
+# if not enabled:
+# return
+#
+# if opt_level not in opt_levels:
+# raise RuntimeError("Unexpected optimization level. Options are 'O0', 'O1', 'O2', 'O3'.")
+# else:
+# opt_properties = opt_levels[opt_level](Properties())
+# print("Selected optimization level {}", opt_levels[opt_level].brief)
+# print("Defaults for this optimization level are:")
+# for k, v in opt_properties.options:
+# print("{:22} : {}".format(k, v))
+#
+# print("Processing user overrides (additional kwargs that are not None)...")
+# for k, v in kwargs:
+# if k not in _amp_state.opt_properties.options:
+# raise RuntimeError("Unexpected kwarg {}".format(k))
+# if v is not None:
+# setattr(opt_properties, k, v)
+#
+# print("After processing overrides, optimization options are:")
+# for k, v in opt_properties.options:
+# print("{:22} : {}".format(k, v))
diff --git a/apex/apex/amp/handle.py b/apex/apex/amp/handle.py
new file mode 100644
index 00000000..0be567ca
--- /dev/null
+++ b/apex/apex/amp/handle.py
@@ -0,0 +1,281 @@
+import contextlib
+import warnings
+import sys
+import torch
+
+from . import utils
+from .opt import OptimWrapper
+from .scaler import LossScaler
+from ._amp_state import _amp_state, master_params, maybe_print
+
+if torch.distributed.is_available():
+ from ..parallel.LARC import LARC
+
+
+# There's no reason to expose the notion of a "handle". Everything can happen through amp.* calls.
+@contextlib.contextmanager
+def scale_loss(loss,
+ optimizers,
+ loss_id=0,
+ model=None,
+ delay_unscale=False,
+ delay_overflow_check=False):
+ """
+ On context manager entrance, creates ``scaled_loss = (loss.float())*current loss scale``.
+ ``scaled_loss`` is yielded so that the user can call ``scaled_loss.backward()``::
+
+ with amp.scale_loss(loss, optimizer) as scaled_loss:
+ scaled_loss.backward()
+
+ On context manager exit (if ``delay_unscale=False``), the gradients are checked for infs/NaNs
+ and unscaled, so that ``optimizer.step()`` can be called.
+
+ .. note::
+ If Amp is using explicit FP32 master params (which is the default for ``opt_level=O2``, and
+ can also be manually enabled by supplying ``master_weights=True`` to ``amp.initialize``)
+ any FP16 gradients are copied to FP32 master gradients before being unscaled.
+ ``optimizer.step()`` will then apply the unscaled master gradients to the master params.
+
+ .. warning::
+ If Amp is using explicit FP32 master params, only the FP32 master gradients will be
+ unscaled. The direct ``.grad`` attributes of any FP16
+ model params will remain scaled after context manager exit.
+ This subtlety affects gradient clipping. See "Gradient clipping" under
+ `Advanced Amp Usage`_ for best practices.
+
+ Args:
+ loss(Tensor): Typically a scalar Tensor. The ``scaled_loss`` that the context
+ manager yields is simply ``loss.float()*loss_scale``, so in principle
+ ``loss`` could have more than one element, as long as you call
+ ``backward()`` on ``scaled_loss`` appropriately within the context manager body.
+ optimizers: All optimizer(s) for which the current backward pass is creating gradients.
+ Must be an optimizer or list of optimizers returned from an earlier call
+ to ``amp.initialize``. For example use with multiple optimizers, see
+ "Multiple models/optimizers/losses" under `Advanced Amp Usage`_.
+ loss_id(int, optional, default=0): When used in conjunction with the ``num_losses`` argument
+ to ``amp.initialize``, enables Amp to use a different loss scale per loss. ``loss_id``
+ must be an integer between 0 and ``num_losses`` that tells Amp which loss is
+ being used for the current backward pass. See "Multiple models/optimizers/losses"
+ under `Advanced Amp Usage`_ for examples. If ``loss_id`` is left unspecified, Amp
+ will use the default global loss scaler for this backward pass.
+ model(torch.nn.Module, optional, default=None): Currently unused, reserved to enable future
+ optimizations.
+ delay_unscale(bool, optional, default=False): ``delay_unscale`` is never necessary, and
+ the default value of ``False`` is strongly recommended.
+ If ``True``, Amp will not unscale the gradients or perform model->master
+ gradient copies on context manager exit.
+ ``delay_unscale=True`` is a minor ninja performance optimization and can result
+ in weird gotchas (especially with multiple models/optimizers/losses),
+ so only use it if you know what you're doing.
+ "Gradient accumulation across iterations" under `Advanced Amp Usage`_
+ illustrates a situation where this CAN (but does not need to) be used.
+
+ .. warning::
+ If ``delay_unscale`` is ``True`` for a given backward pass, ``optimizer.step()`` cannot be
+ called yet after context manager exit, and must wait for another, later backward context
+ manager invocation with ``delay_unscale`` left to False.
+
+ .. _`Advanced Amp Usage`:
+ https://nvidia.github.io/apex/advanced.html
+ """
+ if not hasattr(_amp_state, "opt_properties"):
+ raise RuntimeError("Invoked 'with amp.scale_loss`, but internal Amp state has not been initialized. "
+ "model, optimizer = amp.initialize(model, optimizer, opt_level=...) must be called "
+ "before `with amp.scale_loss`.")
+
+ if not _amp_state.opt_properties.enabled:
+ yield loss
+ return
+
+ if isinstance(optimizers, torch.optim.Optimizer) or ('LARC' in globals() and isinstance(optimizers, LARC)):
+ optimizers = [optimizers]
+
+ loss_scaler = _amp_state.loss_scalers[loss_id]
+ loss_scale = loss_scaler.loss_scale()
+
+ if ((not _amp_state.opt_properties.master_weights)
+ and (not loss_scaler.dynamic)
+ and loss_scale == 1.0):
+ yield loss.float()
+ # Needing to drop the cache here as well is an ugly gotcha.
+ # But for now I think it's necessary to short-circuit.
+ # Probably ok to skip this if not delay_unscale
+ if _amp_state.opt_properties.patch_torch_functions:
+ _amp_state.handle._clear_cache()
+ return
+
+ if not delay_unscale:
+ if isinstance(optimizers, list):
+ for optimizer in optimizers:
+ if not optimizer._amp_stash.params_have_scaled_gradients:
+ optimizer._prepare_amp_backward()
+
+ yield (loss.float())*loss_scale
+
+ if delay_unscale:
+ for optimizer in optimizers:
+ optimizer._amp_stash.params_have_scaled_gradients = True
+ else:
+ # FusedSGD may take care of unscaling as part of their step() methods.
+ # if not isinstance(optimizers, FP16_Optimizer_for_fused):
+ loss_scaler.clear_overflow_state()
+ for optimizer in optimizers:
+ optimizer._post_amp_backward(loss_scaler)
+ optimizer._amp_stash.params_have_scaled_gradients = False
+ # For future fused optimizers that enable sync-free dynamic loss scaling,
+ # should_skip will always be False.
+ should_skip = False if delay_overflow_check else loss_scaler.update_scale()
+ if should_skip:
+ for optimizer in optimizers:
+ if not optimizer._amp_stash.already_patched:
+ # Close on loss_scaler and loss_id as well, to be safe. Probably not
+ # necessary because amp.scale_loss is already creating a temporary scope.
+ def patch_step(opt, loss_scaler, loss_id):
+ opt_step = opt.step
+ def skip_step(closure=None):
+ if closure is not None:
+ raise RuntimeError("Currently, Amp does not support closure use with optimizers.")
+ maybe_print(("Gradient overflow. Skipping step, loss scaler " +
+ "{} reducing loss scale to {}").format(loss_id,
+ loss_scaler.loss_scale()))
+ # TODO: I don't like the special casing for different optimizer implementations.
+ # Maybe skip should delegate to a method owned by the optimizers themselves.
+ if hasattr(opt._amp_stash, "all_fp32_from_fp16_params"):
+ # Clear the master grads that wouldn't be zeroed by model.zero_grad()
+ for param in opt._amp_stash.all_fp32_from_fp16_params:
+ param.grad = None
+ if hasattr(opt, "most_recent_scale"):
+ opt.most_recent_scale = 1.0
+ opt.scale_set_by_backward = False
+ opt.step = opt_step
+ opt._amp_stash.already_patched = False
+ return skip_step
+ optimizer.step = patch_step(optimizer, loss_scaler, loss_id)
+ optimizer._amp_stash.already_patched = True
+
+ # Probably ok to skip this if not delay_unscale
+ if _amp_state.opt_properties.patch_torch_functions:
+ _amp_state.handle._clear_cache()
+
+
+# Free function version of AmpHandle.disable_casts, another step on the
+# path to removing the concept of "AmpHandle"
+@contextlib.contextmanager
+def disable_casts():
+ _amp_state.handle._is_active = False
+ yield
+ _amp_state.handle._is_active = True
+
+
+class AmpHandle(object):
+ def __init__(self, loss_scale="dynamic", enable_caching=True, verbose=False):
+ self._enable_caching = enable_caching
+ self._verbose = verbose
+ self._cache = dict()
+ self._default_scaler = LossScaler(loss_scale)
+ self._is_active = True
+ self._all_wrappers = []
+
+ def is_active(self):
+ return self._is_active
+
+ @contextlib.contextmanager
+ def _disable_casts(self):
+ self._is_active = False
+ yield
+ self._is_active = True
+
+ def wrap_optimizer(self, optimizer, num_loss=1):
+ self._default_scaler = None
+ return OptimWrapper(optimizer, self, num_loss)
+
+ @contextlib.contextmanager
+ def scale_loss(self, loss, optimizer):
+ raise RuntimeError("The old Amp API is no longer supported. Please move to the new API, "
+ "documented here: https://nvidia.github.io/apex/amp.html. Transition guide: "
+ "https://nvidia.github.io/apex/amp.html#transition-guide-for-old-api-users")
+
+ if not self.is_active():
+ yield loss
+ return
+
+ if self._default_scaler is None:
+ raise RuntimeError(
+ 'After calling `handle.wrap_optimizer()`, you must explicitly ' +
+ 'use `optimizer.scale_loss(loss)`.')
+
+ # TODO: this code block is duplicated here and `opt.py`. Unify.
+ loss_scale = self._default_scaler.loss_scale()
+ yield loss * loss_scale
+
+ self._default_scaler.clear_overflow_state()
+ self._default_scaler.unscale(
+ master_params(optimizer),
+ master_params(optimizer),
+ loss_scale)
+ should_skip = self._default_scaler.update_scale()
+ if should_skip:
+ optimizer_step = optimizer.step
+ def skip_step():
+ maybe_print('Gradient overflow, skipping update')
+ optimizer.step = optimizer_step
+ optimizer.step = skip_step
+
+ self._clear_cache()
+
+ def _clear_cache(self):
+ self._cache.clear()
+
+ # Experimental support for saving / restoring uncasted versions of functions
+ def _save_func(self, mod, fn, func):
+ self._all_wrappers.append((mod, fn, func))
+
+ def _deactivate(self):
+ for mod, fn, func in self._all_wrappers:
+ utils.set_func(mod, fn, func)
+ self._all_wrappers = []
+
+ @property
+ def has_cache(self):
+ return self._enable_caching
+
+ @property
+ def cache(self):
+ return self._cache
+
+ def remove_cache(self, param):
+ if self.has_cache and param in self.cache:
+ del self.cache[param]
+
+ @property
+ def verbose(self):
+ return self._verbose
+
+class NoOpHandle(object):
+ def is_active(self):
+ return False
+
+ @contextlib.contextmanager
+ def _disable_casts(self):
+ yield
+
+ def wrap_optimizer(self, optimizer, num_loss=1):
+ return OptimWrapper(optimizer, self, num_loss)
+
+ @contextlib.contextmanager
+ def scale_loss(self, loss, optimizer):
+ yield loss
+
+ @property
+ def has_cache(self):
+ return False
+
+ @property
+ def verbose(self):
+ return False
+
+ def _clear_cache(self):
+ pass
+
+ def _deactivate(self):
+ pass
diff --git a/tools/convert_checkpoint/custom/llavaov_1_5/__init__.py b/apex/apex/amp/lists/__init__.py
similarity index 100%
rename from tools/convert_checkpoint/custom/llavaov_1_5/__init__.py
rename to apex/apex/amp/lists/__init__.py
diff --git a/apex/apex/amp/lists/functional_overrides.py b/apex/apex/amp/lists/functional_overrides.py
new file mode 100644
index 00000000..dd009cec
--- /dev/null
+++ b/apex/apex/amp/lists/functional_overrides.py
@@ -0,0 +1,80 @@
+
+# TODO: think about the following two. They do weird things.
+# - torch.nn.utils.clip_grad (but it should always be fp32 anyway)
+# - torch.nn.utils.weight_norm
+
+# Notes:
+# F.instance_norm uses batch_norm internally. Which correctly handles
+# fp16 in/out with fp32 weights. So we shouldn't do anything for
+# either of these.
+# F.normalize calls `input.norm()` internally, so it's redundant, but
+# kept here in case impl. changes.
+# F.cosine_similarity is same: calls `x.norm()` internally.
+
+import torch.nn.functional
+
+MODULE = torch.nn.functional
+
+FP16_FUNCS = [
+ 'conv1d',
+ 'conv2d',
+ 'conv3d',
+ 'conv_transpose1d',
+ 'conv_transpose2d',
+ 'conv_transpose3d',
+ 'conv_tbc', # Undocumented / maybe new?
+ 'linear',
+]
+
+FP32_FUNCS = [
+
+ # Interpolation/Upsampling TODO: Remove for 1.2
+ 'interpolate',
+ 'grid_sample',
+
+ # Pointwise
+ 'softplus',
+ 'softmin',
+ 'log_softmax',
+ 'softmax',
+ 'gelu',
+
+ # Normalization
+ 'layer_norm',
+ 'group_norm',
+ 'local_response_norm',
+ 'normalize',
+ 'cosine_similarity',
+
+ # Loss functions
+ # TODO: which of these can be fp16?
+ 'poisson_nll_loss',
+ 'cosine_embedding_loss',
+ 'cross_entropy',
+ 'hinge_embedding_loss',
+ 'kl_div',
+ 'l1_loss',
+ 'mse_loss',
+ 'margin_ranking_loss',
+ 'multilabel_margin_loss',
+ 'multilabel_soft_margin_loss',
+ 'multi_margin_loss',
+ 'nll_loss',
+ 'binary_cross_entropy_with_logits',
+ 'smooth_l1_loss',
+ 'soft_margin_loss',
+ 'triplet_margin_loss',
+ 'ctc_loss'
+]
+
+BANNED_FUNCS = [
+ ('binary_cross_entropy',
+ ("\namp does not work out-of-the-box with `F.binary_cross_entropy` or `torch.nn.BCELoss.` "
+ "It requires that the output of the previous function be already a FloatTensor. \n\n"
+ "Most models have a Sigmoid right before BCELoss. In that case, you can use\n"
+ " torch.nn.BCEWithLogitsLoss\nto combine Sigmoid+BCELoss into a single layer "
+ "that is compatible with amp.\nAnother option is to add\n"
+ " amp.register_float_function(torch, 'sigmoid')\nbefore calling `amp.init()`.\n"
+ "If you _really_ know what you are doing, you can disable this warning by passing "
+ "allow_banned=True to `amp.init()`."))
+]
diff --git a/apex/apex/amp/lists/tensor_overrides.py b/apex/apex/amp/lists/tensor_overrides.py
new file mode 100644
index 00000000..18f3e5dc
--- /dev/null
+++ b/apex/apex/amp/lists/tensor_overrides.py
@@ -0,0 +1,63 @@
+from .. import compat
+from . import torch_overrides
+
+import importlib
+
+import torch
+
+# if compat.variable_is_tensor() and not compat.tensor_is_variable():
+MODULE = torch.Tensor
+# else:
+# MODULE = torch.autograd.Variable
+
+
+FP16_FUNCS = compat.filter_attrs(MODULE, [
+ '__matmul__',
+])
+
+FP32_FUNCS = compat.filter_attrs(MODULE, [
+ '__ipow__',
+ '__pow__',
+ '__rpow__',
+
+ # Cast to fp32 before transfer to CPU
+ 'cpu',
+])
+
+CASTS = compat.filter_attrs(MODULE, [
+ '__add__',
+ '__div__',
+ '__eq__',
+ '__ge__',
+ '__gt__',
+ '__iadd__',
+ '__idiv__',
+ '__imul__',
+ '__isub__',
+ '__itruediv__',
+ '__le__',
+ '__lt__',
+ '__mul__',
+ '__ne__',
+ '__radd__',
+ '__rdiv__',
+ '__rmul__',
+ '__rsub__',
+ '__rtruediv__',
+ '__sub__',
+ '__truediv__',
+])
+
+# None of these, but here to make code cleaner.
+SEQUENCE_CASTS = []
+
+# We need to grab all the methods from torch_overrides and add them to
+# the Tensor lists as well, as almost all methods are duplicated
+# between `torch` and `torch.Tensor` (and check with `hasattr`,
+# because a few random ones aren't defined on Tensor)
+_self_mod = importlib.import_module(__name__)
+for attrname in ['FP16_FUNCS', 'FP32_FUNCS', 'CASTS', 'SEQUENCE_CASTS']:
+ lst = getattr(_self_mod, attrname)
+ for fn in getattr(torch_overrides, attrname):
+ if hasattr(MODULE, fn):
+ lst.append(fn)
diff --git a/apex/apex/amp/lists/torch_overrides.py b/apex/apex/amp/lists/torch_overrides.py
new file mode 100644
index 00000000..7dedb05a
--- /dev/null
+++ b/apex/apex/amp/lists/torch_overrides.py
@@ -0,0 +1,115 @@
+import torch
+
+from .. import utils
+
+MODULE = torch
+
+FP16_FUNCS = [
+ # Low level functions wrapped by torch.nn layers.
+ # The wrapper layers contain the weights which are then passed in as a parameter
+ # to these functions.
+ 'conv1d',
+ 'conv2d',
+ 'conv3d',
+ 'conv_transpose1d',
+ 'conv_transpose2d',
+ 'conv_transpose3d',
+ 'conv_tbc',
+ 'prelu',
+
+ # BLAS
+ 'addmm',
+ 'addmv',
+ 'addr',
+ 'matmul',
+ 'mm',
+ 'mv',
+]
+
+FP32_FUNCS = [
+ # Pointwise
+ 'acos',
+ 'asin',
+ 'cosh',
+ 'erfinv',
+ 'exp',
+ 'expm1',
+ 'log',
+ 'log10',
+ 'log2',
+ 'reciprocal',
+ 'rsqrt',
+ 'sinh',
+ 'tan',
+
+ # Other math
+ 'pow',
+
+ # Reduction
+ 'cumprod',
+ 'cumsum',
+ 'dist',
+ # 'mean',
+ 'norm',
+ 'prod',
+ 'std',
+ 'sum',
+ 'var',
+
+ # Misc
+ 'renorm'
+]
+
+version_strings = torch.__version__.split('.')
+version_major = version_strings[0]
+version_minor = version_strings[1]
+version_num = float(version_major + "." + version_minor)
+# Before torch 1.1, mean must be blacklisted.
+if version_num < 1.1:
+ FP32_FUNCS.append('mean')
+
+# Before CUDA 9.1, batched matmul was missing fast FP16 kernels. We
+# check the CUDA version -- if at least 9.1, then put the bmm
+# functions on the fp16 list. Otherwise, put them on the fp32 list.
+_bmms = ['addbmm',
+ 'baddbmm',
+ 'bmm']
+
+if utils.is_cuda_enabled():
+ # workaround https://github.com/facebookresearch/maskrcnn-benchmark/issues/802
+ if utils.get_cuda_version() >= (9, 1, 0):
+ FP16_FUNCS.extend(_bmms)
+ else:
+ FP32_FUNCS.extend(_bmms)
+
+# Multi-tensor fns that may need type promotion
+CASTS = [
+ # Multi-tensor math
+ 'addcdiv',
+ 'addcmul',
+ 'atan2',
+ 'cross',
+ 'bilinear',
+ 'dot',
+
+ # Element-wise _or_ tensor-wise math
+ 'add',
+ 'div',
+ 'mul',
+
+ # Comparison
+ 'eq',
+ 'equal',
+ 'ge',
+ 'gt',
+ 'le',
+ 'lt',
+ 'ne'
+]
+
+# Functions that take sequence arguments. We need to inspect the whole
+# sequence and cast to the widest type.
+SEQUENCE_CASTS = [
+ 'cat',
+ 'stack'
+]
diff --git a/apex/apex/amp/opt.py b/apex/apex/amp/opt.py
new file mode 100644
index 00000000..baf31168
--- /dev/null
+++ b/apex/apex/amp/opt.py
@@ -0,0 +1,103 @@
+import contextlib
+import warnings
+
+from .scaler import LossScaler, master_params
+from ._amp_state import maybe_print
+
+import numpy as np
+
+class OptimWrapper(object):
+ def __init__(self, optimizer, amp_handle, num_loss):
+ self._optimizer = optimizer
+ self._amp_handle = amp_handle
+ self._num_loss = num_loss
+ self._loss_idx = 0
+ self._skip_next = [False] * num_loss
+ self._loss_scaler = [LossScaler('dynamic') for _ in range(num_loss)]
+
+ @contextlib.contextmanager
+ def scale_loss(self, loss):
+ if not self._amp_handle.is_active():
+ yield loss
+ return
+
+ # When there are multiple losses per-optimizer, we need
+ # to save out current grad accumulation, since we won't be
+ # able to unscale this particulare loss once the grads are
+ # all mixed together.
+ cached_grads = []
+ if self._loss_idx > 0:
+ for p in master_params(self._optimizer):
+ if p.grad is not None:
+ cached_grads.append(p.grad.data.detach().clone())
+ else:
+ cached_grads.append(None)
+ self._optimizer.zero_grad()
+
+ loss_scale = self._cur_loss_scaler().loss_scale()
+ yield loss * loss_scale
+
+ self._cur_loss_scaler().clear_overflow_state()
+ self._cur_loss_scaler().unscale(
+ master_params(self._optimizer),
+ master_params(self._optimizer),
+ loss_scale)
+ self._skip_next[self._loss_idx] = self._cur_loss_scaler().update_scale()
+ self._loss_idx += 1
+
+ if len(cached_grads) > 0:
+ for p, cached_grad in zip(master_params(self._optimizer),
+ cached_grads):
+ if cached_grad is not None:
+ p.grad.data.add_(cached_grad)
+ cached_grads = []
+
+ def _cur_loss_scaler(self):
+ assert 0 <= self._loss_idx < self._num_loss
+ return self._loss_scaler[self._loss_idx]
+
+ def step(self, closure=None):
+ if not self._amp_handle.is_active():
+ return self._optimizer.step(closure=closure)
+
+ self._loss_idx = 0
+
+ for group in self._optimizer.param_groups:
+ for p in group['params']:
+ self._amp_handle.remove_cache(p)
+
+ if closure is not None:
+ raise NotImplementedError(
+ 'The `closure` argument is unsupported by the amp ' +
+ 'optimizer wrapper.')
+ if any(self._skip_next):
+ maybe_print('Gradient overflow, skipping update')
+ self._skip_next = [False] * self._num_loss
+ else:
+ return self._optimizer.step(closure=closure)
+
+ # Forward any attribute lookups
+ def __getattr__(self, attr):
+ return getattr(self._optimizer, attr)
+
+ # Forward all torch.optim.Optimizer methods
+ def __getstate__(self):
+ return self._optimizer.__getstate__()
+
+ def __setstate__(self):
+ return self._optimizer.__setstate__()
+
+ def __repr__(self):
+ return self._optimizer.__repr__()
+
+ def state_dict(self):
+ return self._optimizer.state_dict()
+
+ def load_state_dict(self, state_dict):
+ return self._optimizer.load_state_dict(state_dict)
+
+ def zero_grad(self):
+ return self._optimizer.zero_grad()
+
+ def add_param_group(self, param_group):
+ return self._optimizer.add_param_group(param_group)
diff --git a/apex/apex/amp/rnn_compat.py b/apex/apex/amp/rnn_compat.py
new file mode 100644
index 00000000..d062ae26
--- /dev/null
+++ b/apex/apex/amp/rnn_compat.py
@@ -0,0 +1,53 @@
+from . import utils, wrap
+
+import torch
+_VF = torch._C._VariableFunctions
+RNN_NAMES = ['rnn_relu', 'rnn_tanh', 'gru', 'lstm']
+
+def _gen_VF_wrapper(name):
+ def wrapper(*args, **kwargs):
+ return getattr(_VF, name)(*args, **kwargs)
+ return wrapper
+
+# Some python magic to generate an object that has the rnn cell functions
+# defined on it, all of which call into corresponding _VF version.
+# Intended to patch torch.nn.modules.rnn._VF (aka, the ref named "_VF"
+# imported at module scope within torch.nn.modules.rnn). This should
+# not affect third-party importers of _VF.py.
+class VariableFunctionsShim(object):
+ def __init__(self):
+ for name in RNN_NAMES:
+ for suffix in ['', '_cell']:
+ fn_name = name + suffix
+ setattr(self, fn_name, _gen_VF_wrapper(fn_name))
+
+def has_old_rnns():
+ try:
+ torch.nn.backends.thnn.backend.LSTMCell
+ return True
+ except:
+ return False
+
+def whitelist_rnn_cells(handle, verbose):
+ # Different module + function names in old/new RNN cases
+ if has_old_rnns():
+ fn_names = ['RNNReLUCell', 'RNNTanhCell', 'LSTMCell', 'GRUCell']
+ mod = torch.nn.backends.thnn.backend
+ else:
+ fn_names = [x + '_cell' for x in RNN_NAMES]
+ mod = torch.nn.modules.rnn._VF
+ assert isinstance(mod, VariableFunctionsShim)
+
+ # Insert casts on cell functions
+ for fn in fn_names:
+ wrap.cached_cast(mod, fn, utils.maybe_half, handle,
+ try_caching=True, verbose=verbose)
+
+ if has_old_rnns():
+ # Special handling of `backward` for fused gru / lstm:
+ # The `backward` method calls Tensor.sum() (blacklist) internally,
+ # and then the resulting grad_input has the wrong type.
+ # TODO: where else is this a problem?
+ for rnn_type in ['GRUFused', 'LSTMFused']:
+ mod = getattr(torch.nn._functions.thnn.rnnFusedPointwise, rnn_type)
+ wrap.disable_casts(mod, 'backward', handle)
diff --git a/apex/apex/amp/scaler.py b/apex/apex/amp/scaler.py
new file mode 100644
index 00000000..99888bc6
--- /dev/null
+++ b/apex/apex/amp/scaler.py
@@ -0,0 +1,217 @@
+import torch
+from ..multi_tensor_apply import multi_tensor_applier
+from ._amp_state import _amp_state, master_params, maybe_print
+from itertools import product
+
+def scale_check_overflow_python(model_grad, master_grad, scale, check_overflow=False):
+ # Exception handling for 18.04 compatibility
+ if check_overflow:
+ cpu_sum = float(model_grad.float().sum())
+ if cpu_sum == float('inf') or cpu_sum == -float('inf') or cpu_sum != cpu_sum:
+ return True
+
+ if master_grad is not model_grad: # copy_ probably internally short-circuits this
+ master_grad.copy_(model_grad)
+ if scale != 1.0:
+ master_grad.mul_(scale)
+ return False
+
+def axpby_check_overflow_python(model_grad, stashed_grad, master_grad, a, b, check_overflow=False):
+ # Exception handling for 18.04 compatibility
+ if check_overflow:
+ cpu_sum = float(model_grad.float().sum())
+ if cpu_sum == float('inf') or cpu_sum == -float('inf') or cpu_sum != cpu_sum:
+ return True
+
+ # if master_grad is not model_grad: # copy_ probably internally short-circuits this
+ # master_grad.copy_(model_grad)
+ assert stashed_grad.dtype == master_grad.dtype
+ converted_model_grad = model_grad.data.to(master_grad.dtype)
+ master_grad.data = a*converted_model_grad.data + b*stashed_grad.data
+ return False
+
+class LossScaler(object):
+ warned_no_fused_kernel = False
+ warned_unscaling_non_fp32_grad = False
+ has_fused_kernel = False
+
+ def __init__(self,
+ loss_scale,
+ init_scale=2.**16,
+ scale_factor=2.,
+ scale_window=2000,
+ min_loss_scale=None,
+ max_loss_scale=2.**24):
+ if loss_scale == "dynamic":
+ self.dynamic = True
+ self._loss_scale = min(max_loss_scale, init_scale)
+ else:
+ self.dynamic = False
+ self._loss_scale = loss_scale
+ self._max_loss_scale = max_loss_scale
+ self._min_loss_scale = min_loss_scale
+ self._scale_seq_len = scale_window
+ self._unskipped = 0
+ self._has_overflow = False
+ self._overflow_buf = torch.cuda.IntTensor([0])
+ if multi_tensor_applier.available:
+ import amp_C
+ LossScaler.has_fused_kernel = multi_tensor_applier.available
+ LossScaler.multi_tensor_scale_cuda = amp_C.multi_tensor_scale
+ LossScaler.multi_tensor_axpby_cuda = amp_C.multi_tensor_axpby
+ else:
+ if not LossScaler.warned_no_fused_kernel:
+ maybe_print(
+ "Warning: multi_tensor_applier fused unscale kernel is unavailable, "
+ "possibly because apex was installed without --cuda_ext --cpp_ext. "
+ "Using Python fallback. Original ImportError was: " +
+ repr(multi_tensor_applier.import_err),
+ True)
+ LossScaler.has_fused_kernel = False
+ LossScaler.warned_no_fused_kernel = True
+
+ def loss_scale(self):
+ return self._loss_scale
+
+ def unscale_python(self, model_grads, master_grads, scale):
+ for model, master in zip(model_grads, master_grads):
+ if model is not None:
+ if not LossScaler.warned_unscaling_non_fp32_grad:
+ if master.dtype != torch.float32:
+ maybe_print(
+ "Attempting to unscale a grad with type {} ".format(master.type()) +
+ "Unscaling non-fp32 grads may indicate an error. "
+ "When using Amp, you don't need to call .half() on your model.")
+ LossScaler.warned_unscaling_non_fp32_grad = True
+ self._has_overflow = scale_check_overflow_python(model,
+ master,
+ 1./scale,
+ self.dynamic)
+ if self._has_overflow and self.dynamic:
+ break
+
+ # unused_scale keeps some of the old API alive for hopefully a short time.
+ def unscale(self, model_grads, master_grads, unused_scale, models_are_masters=False, scale_override=None):
+ if self._has_overflow:
+ return
+
+ scale = self._loss_scale
+ if scale_override is not None:
+ scale = scale_override
+
+ if scale == 1.0 and models_are_masters and not self.dynamic:
+ return
+
+ if LossScaler.has_fused_kernel:
+ # if (not LossScaler.warned_unscaling_non_fp32_grad
+ # and master_grads[0].dtype == torch.float16):
+ # print("Warning: unscaling grads that are not FP32. "
+ # "Unscaling non-fp32 grads may indicate an error. "
+ # "When using Amp, you don't need to call .half() on your model.")
+ # # Setting this to True unconditionally allows the possibility of an escape
+ # # if never-before-seen non-fp32 grads are created in some later iteration.
+ # LossScaler.warned_unscaling_non_fp32_grad = True
+ multi_tensor_applier(LossScaler.multi_tensor_scale_cuda,
+ self._overflow_buf,
+ [model_grads, master_grads],
+ 1./scale)
+ else:
+ self.unscale_python(model_grads, master_grads, scale)
+
+ # Defer to update_scale
+ # If the fused kernel is available, we only need one D2H memcopy and sync.
+ # if LossScaler.has_fused_kernel and self.dynamic and not self._has_overflow:
+ # self._has_overflow = self._overflow_buf.item()
+
+ def unscale_with_stashed_python(self,
+ model_grads,
+ stashed_master_grads,
+ master_grads,
+ a,
+ b):
+ for model, stashed, master in zip(model_grads, stashed_master_grads, master_grads):
+ if model is None and stashed is None:
+ continue
+ else:
+ if not LossScaler.warned_unscaling_non_fp32_grad:
+ if master.dtype != torch.float32:
+ maybe_print(
+ "Attempting to unscale a grad with type {} ".format(master.type()) +
+ "Unscaling non-fp32 grads may indicate an error. "
+ "When using Amp, you don't need to call .half() on your model.")
+ LossScaler.warned_unscaling_non_fp32_grad = True
+ self._has_overflow = axpby_check_overflow_python(model,
+ stashed,
+ master,
+ a,
+ b,
+ self.dynamic)
+ if self._has_overflow and self.dynamic:
+ break
+
+ def unscale_with_stashed(self,
+ model_grads,
+ stashed_master_grads,
+ master_grads,
+ scale_override=None):
+ if self._has_overflow:
+ return
+
+ grads_have_scale, stashed_have_scale, out_scale = self._loss_scale, 1.0, 1.0
+ if scale_override is not None:
+ grads_have_scale, stashed_have_scale, out_scale = scale_override
+
+ if LossScaler.has_fused_kernel:
+ if (not LossScaler.warned_unscaling_non_fp32_grad
+ and master_grads[0].dtype == torch.float16):
+ print("Warning: unscaling grads that are not FP32. "
+ "Unscaling non-fp32 grads may indicate an error. "
+ "When using Amp, you don't need to call .half() on your model.")
+ # Setting this to True unconditionally allows the possibility of an escape
+ # if never-before-seen non-fp32 grads are created in some later iteration.
+ LossScaler.warned_unscaling_non_fp32_grad = True
+ multi_tensor_applier(LossScaler.multi_tensor_axpby_cuda,
+ self._overflow_buf,
+ [model_grads, stashed_master_grads, master_grads],
+ out_scale/grads_have_scale, # 1./scale,
+ out_scale/stashed_have_scale, # 1.0,
+ 0) # check only arg 0, aka the incoming model grads, for infs
+ else:
+ self.unscale_with_stashed_python(model_grads,
+ stashed_master_grads,
+ master_grads,
+ out_scale/grads_have_scale,
+ out_scale/stashed_have_scale)
+
+ # Defer to update_scale
+ # If the fused kernel is available, we only need one D2H memcopy and sync.
+ # if LossScaler.has_fused_kernel and self.dynamic and not self._has_overflow:
+ # self._has_overflow = self._overflow_buf.item()
+
+ def clear_overflow_state(self):
+ self._has_overflow = False
+ if self.has_fused_kernel:
+ self._overflow_buf.zero_()
+
+ # Separate so unscale() can be called more that once before updating.
+ def update_scale(self):
+ # If the fused kernel is available, we only need one D2H memcopy and sync.
+ if LossScaler.has_fused_kernel and self.dynamic and not self._has_overflow:
+ self._has_overflow = self._overflow_buf.item()
+
+ if self._has_overflow and self.dynamic:
+ should_skip = True
+ if(self._min_loss_scale):
+ self._loss_scale = max(self._min_loss_scale, self._loss_scale/2.)
+ else:
+ self._loss_scale = self._loss_scale/2.
+ self._unskipped = 0
+ else:
+ should_skip = False
+ self._unskipped += 1
+
+ if self._unskipped == self._scale_seq_len and self.dynamic:
+ self._loss_scale = min(self._max_loss_scale, self._loss_scale*2.)
+ self._unskipped = 0
+
+ return should_skip
diff --git a/apex/apex/amp/utils.py b/apex/apex/amp/utils.py
new file mode 100644
index 00000000..0590cd70
--- /dev/null
+++ b/apex/apex/amp/utils.py
@@ -0,0 +1,210 @@
+from . import compat
+
+import functools
+import itertools
+
+import torch
+
+def is_cuda_enabled():
+ return torch.version.cuda is not None
+
+def get_cuda_version():
+ return tuple(int(x) for x in torch.version.cuda.split('.'))
+
+def is_fp_tensor(x):
+ if is_nested(x):
+ # Fast-fail version of all(is_fp_tensor)
+ for y in x:
+ if not is_fp_tensor(y):
+ return False
+ return True
+ return compat.is_tensor_like(x) and compat.is_floating_point(x)
+
+def is_nested(x):
+ return isinstance(x, tuple) or isinstance(x, list)
+
+def should_cache(x):
+ if is_nested(x):
+ # Fast-fail version of all(should_cache)
+ for y in x:
+ if not should_cache(y):
+ return False
+ return True
+ return isinstance(x, torch.nn.parameter.Parameter) and \
+ type_string(x) == 'FloatTensor'
+
+def collect_fp_tensor_types(args, kwargs):
+ def collect_types(x, types):
+ if is_nested(x):
+ for y in x:
+ collect_types(y, types)
+ else:
+ types.add(type_string(x))
+
+ all_args = itertools.chain(args, kwargs.values())
+ types = set()
+ for x in all_args:
+ if is_fp_tensor(x):
+ collect_types(x, types)
+ return types
+
+def type_string(x):
+ return x.type().split('.')[-1]
+
+def maybe_half(x, name='', verbose=False):
+ if is_nested(x):
+ return type(x)([maybe_half(y) for y in x])
+
+ if not x.is_cuda or type_string(x) == 'HalfTensor':
+ return x
+ else:
+ if verbose:
+ print('Float->Half ({})'.format(name))
+ return x.half()
+
+def maybe_float(x, name='', verbose=False):
+ if is_nested(x):
+ return type(x)([maybe_float(y) for y in x])
+
+ if not x.is_cuda or type_string(x) == 'FloatTensor':
+ return x
+ else:
+ if verbose:
+ print('Half->Float ({})'.format(name))
+ return x.float()
+
+# NB: returneds casted `args`, mutates `kwargs` in-place
+def casted_args(cast_fn, args, kwargs):
+ new_args = []
+ for x in args:
+ if is_fp_tensor(x):
+ new_args.append(cast_fn(x))
+ else:
+ new_args.append(x)
+ for k in kwargs:
+ val = kwargs[k]
+ if is_fp_tensor(val):
+ kwargs[k] = cast_fn(val)
+ return new_args
+
+def cached_cast(cast_fn, x, cache):
+ if is_nested(x):
+ return type(x)([cached_cast(y) for y in x])
+ if x in cache:
+ cached_x = cache[x]
+ if x.requires_grad and cached_x.requires_grad:
+ # Make sure x is actually cached_x's autograd parent.
+ if cached_x.grad_fn.next_functions[1][0].variable is not x:
+ raise RuntimeError("x and cache[x] both require grad, but x is not "
+ "cache[x]'s parent. This is likely an error.")
+ # During eval, it's possible to end up caching casted weights with
+ # requires_grad=False. On the next training iter, if cached_x is found
+ # and reused from the cache, it will not actually have x as its parent.
+ # Therefore, we choose to invalidate the cache (and force refreshing the cast)
+ # if x.requires_grad and cached_x.requires_grad do not match.
+ #
+ # During eval (i.e. running under with torch.no_grad()) the invalidation
+ # check would cause the cached value to be dropped every time, because
+ # cached_x would always be created with requires_grad=False, while x would
+ # still have requires_grad=True. This would render the cache effectively
+ # useless during eval. Therefore, if we are running under the no_grad()
+ # context manager (torch.is_grad_enabled=False) we elide the invalidation
+ # check, and use the cached value even though its requires_grad flag doesn't
+ # match. During eval, we don't care that there's no autograd-graph
+ # connection between x and cached_x.
+ if torch.is_grad_enabled() and x.requires_grad != cached_x.requires_grad:
+ del cache[x]
+ else:
+ return cached_x
+
+ casted_x = cast_fn(x)
+ cache[x] = casted_x
+ return casted_x
+
+def verbosify(cast_fn, fn_name, verbose):
+ if verbose:
+ return functools.partial(cast_fn, name=fn_name, verbose=verbose)
+ else:
+ return cast_fn
+
+def as_inplace(fns):
+ for x in fns:
+ yield x + '_'
+
+def has_func(mod, fn):
+ if isinstance(mod, dict):
+ return fn in mod
+ else:
+ return hasattr(mod, fn)
+
+def get_func(mod, fn):
+ if isinstance(mod, dict):
+ return mod[fn]
+ else:
+ return getattr(mod, fn)
+
+def set_func(mod, fn, new_fn):
+ if isinstance(mod, dict):
+ mod[fn] = new_fn
+ else:
+ setattr(mod, fn, new_fn)
+
+def set_func_save(handle, mod, fn, new_fn):
+ cur_fn = get_func(mod, fn)
+ handle._save_func(mod, fn, cur_fn)
+ set_func(mod, fn, new_fn)
+
+# A couple problems get solved here:
+# - The flat_weight buffer is disconnected from autograd graph,
+# so the fp16 weights need to be derived from the input weights
+# to this forward call, not the flat buffer.
+# - The ordering of weights in the flat buffer is...idiosyncratic.
+# First problem is solved with combination of set_ (to set up
+# correct storage) and copy_ (so the fp16 weight derives from the
+# fp32 one in autograd.
+# Second is solved by doing ptr arithmetic on the fp32 weights
+# to derive the correct offset.
+#
+# TODO: maybe this should actually use
+# `torch._cudnn_rnn_flatten_weight`? But then I need to call
+# on first iter and cache the right offsets. Ugh.
+def synthesize_flattened_rnn_weights(fp32_weights,
+ fp16_flat_tensor,
+ rnn_fn='',
+ verbose=False):
+ fp16_weights = []
+ fp32_base_ptr = fp32_weights[0][0].data_ptr()
+ for layer_weights in fp32_weights:
+ fp16_layer_weights = []
+ for w_fp32 in layer_weights:
+ w_fp16 = w_fp32.new().half()
+ offset = (w_fp32.data_ptr() - fp32_base_ptr) // w_fp32.element_size()
+ w_fp16.set_(fp16_flat_tensor.storage(),
+ offset,
+ w_fp32.shape)
+ w_fp16.copy_(w_fp32)
+ if verbose:
+ print('Float->Half ({})'.format(rnn_fn))
+ fp16_layer_weights.append(w_fp16)
+ fp16_weights.append(fp16_layer_weights)
+ return fp16_weights
+
+# Roughly same as above, just the `fp32_weights` aren't nested.
+# Code kept separate for readability.
+def new_synthesize_flattened_rnn_weights(fp32_weights,
+ fp16_flat_tensor,
+ rnn_fn='',
+ verbose=False):
+ fp16_weights = []
+ fp32_base_ptr = fp32_weights[0].data_ptr()
+ for w_fp32 in fp32_weights:
+ w_fp16 = w_fp32.new().half()
+ offset = (w_fp32.data_ptr() - fp32_base_ptr) // w_fp32.element_size()
+ w_fp16.set_(fp16_flat_tensor.storage(),
+ offset,
+ w_fp32.shape)
+ w_fp16.copy_(w_fp32)
+ if verbose:
+ print('Float->Half ({})'.format(rnn_fn))
+ fp16_weights.append(w_fp16)
+ return fp16_weights
diff --git a/apex/apex/amp/wrap.py b/apex/apex/amp/wrap.py
new file mode 100644
index 00000000..559d0558
--- /dev/null
+++ b/apex/apex/amp/wrap.py
@@ -0,0 +1,276 @@
+from . import compat
+from . import utils
+from ._amp_state import _amp_state
+from . import rnn_compat
+
+import functools
+
+import torch
+
+def make_cast_wrapper(orig_fn, cast_fn, handle,
+ try_caching=False):
+ @functools.wraps(orig_fn)
+ def wrapper(*args, **kwargs):
+ if not handle.is_active():
+ return orig_fn(*args, **kwargs)
+
+ if try_caching and handle.has_cache:
+ args = list(args)
+ for i in range(len(args)):
+ if utils.should_cache(args[i]):
+ args[i] = utils.cached_cast(cast_fn, args[i], handle.cache)
+ for k in kwargs:
+ if utils.should_cache(kwargs[k]):
+ kwargs[k] = utils.cached_cast(cast_fn, kwargs[k], handle.cache)
+ new_args = utils.casted_args(cast_fn,
+ args,
+ kwargs)
+ return orig_fn(*new_args, **kwargs)
+ return wrapper
+
+def cached_cast(mod, fn, cast_fn, handle,
+ try_caching=False, verbose=False):
+ if not utils.has_func(mod, fn):
+ return
+
+ orig_fn = utils.get_func(mod, fn)
+ cast_fn = utils.verbosify(cast_fn, fn, verbose)
+ wrapper = make_cast_wrapper(orig_fn, cast_fn, handle, try_caching)
+ utils.set_func_save(handle, mod, fn, wrapper)
+
+# `handle` arg is unused, but simplifies API to make `make_cast_wrapper`
+# Annoyingly, make_promote_wrapper still uses the global handle. Once everyone
+# is on the new API and I am free to get rid of handle, I can clean this up.
+def make_promote_wrapper(orig_fn, cast_fn, handle=None):
+ @functools.wraps(orig_fn)
+ def wrapper(*args, **kwargs):
+ if not _amp_state.handle.is_active():
+ return orig_fn(*args, **kwargs)
+
+ types = utils.collect_fp_tensor_types(args, kwargs)
+
+ if len(types) <= 1:
+ return orig_fn(*args, **kwargs)
+ elif len(types) == 2 and types == set(['HalfTensor', 'FloatTensor']):
+ new_args = utils.casted_args(cast_fn,
+ args,
+ kwargs)
+ return orig_fn(*new_args, **kwargs)
+ else:
+ raise NotImplementedError('Do not know how to handle ' +
+ 'these types to promote: {}'
+ .format(types))
+ return wrapper
+
+def promote(mod, fn, handle, verbose=False):
+ orig_fn = utils.get_func(mod, fn)
+ maybe_float = utils.verbosify(utils.maybe_float, fn, verbose)
+ wrapper = make_promote_wrapper(orig_fn, maybe_float)
+ utils.set_func_save(handle, mod, fn, wrapper)
+
+def sequence_promote(mod, fn, handle, verbose=False):
+ orig_fn = utils.get_func(mod, fn)
+ maybe_float = utils.verbosify(utils.maybe_float, fn, verbose)
+ @functools.wraps(orig_fn)
+ def wrapper(seq, *args, **kwargs):
+ if not _amp_state.handle.is_active():
+ return orig_fn(seq, *args, **kwargs)
+
+ types = set([utils.type_string(x) for x in seq])
+ if len(types) <= 1:
+ return orig_fn(seq, *args, **kwargs)
+ elif types == set(['HalfTensor', 'FloatTensor']):
+ cast_seq = utils.casted_args(maybe_float,
+ seq, {})
+ return orig_fn(cast_seq, *args, **kwargs)
+ else:
+ # TODO: other mixed-type cases aren't due to amp.
+ # Just pass through?
+ return orig_fn(seq, *args, **kwargs)
+ utils.set_func_save(handle, mod, fn, wrapper)
+
+def promote_match_arg0(mod, fn, handle, verbose=False):
+ if not utils.has_func(mod, fn):
+ return
+
+ orig_fn = utils.get_func(mod, fn)
+ @functools.wraps(orig_fn)
+ def wrapper(arg0, *args, **kwargs):
+ assert compat.is_tensor_like(arg0)
+ if not _amp_state.handle.is_active():
+ return orig_fn(arg0, *args, **kwargs)
+
+ if utils.type_string(arg0) == 'HalfTensor':
+ cast_fn = utils.maybe_half
+ elif utils.type_string(arg0) == 'FloatTensor':
+ cast_fn = utils.maybe_float
+ else:
+ return orig_fn(arg0, *args, **kwargs)
+ cast_fn = utils.verbosify(cast_fn, fn, verbose)
+ new_args = utils.casted_args(cast_fn, args, kwargs)
+ return orig_fn(arg0, *new_args, **kwargs)
+ utils.set_func_save(handle, mod, fn, wrapper)
+
+def err_if_any_half(mod, fn, handle, custom_err_msg=None):
+ if not utils.has_func(mod, fn):
+ return
+
+ orig_fn = utils.get_func(mod, fn)
+ @functools.wraps(orig_fn)
+ def wrapper(*args, **kwargs):
+ types = utils.collect_fp_tensor_types(args, kwargs)
+ if 'HalfTensor' in types:
+ if custom_err_msg:
+ raise NotImplementedError(custom_err_msg)
+ else:
+ raise NotImplementedError('Cannot call in-place function ' +
+ '{} with fp16 arguments.'.format(fn))
+ else:
+ return orig_fn(*args, **kwargs)
+ utils.set_func_save(handle, mod, fn, wrapper)
+
+def err_if_arg0_half(mod, fn, handle, verbose=False):
+ if not utils.has_func(mod, fn):
+ return
+
+ orig_fn = utils.get_func(mod, fn)
+ @functools.wraps(orig_fn)
+ def wrapper(arg0, *args, **kwargs):
+ assert compat.is_tensor_like(arg0)
+ if utils.type_string(arg0) == 'HalfTensor':
+ raise NotImplementedError('Cannot call in-place method ' +
+ '{} on fp16 Tensors.'.format(fn))
+ else:
+ cast_fn = utils.verbosify(utils.maybe_float, fn, verbose)
+ new_args = utils.casted_args(cast_fn, args, kwargs)
+ return orig_fn(arg0, *new_args, **kwargs)
+ utils.set_func_save(handle, mod, fn, wrapper)
+
+# Current RNN approach:
+# - Wrap top-level `RNN` function in thnn backend
+# - Will call into either CudnnRNN or AutogradRNN
+# - Each of these are factory functions that return a per-iter
+# `forward` function
+# - We interpose on the factory function to:
+# 1) Interpose on the actual forward function and put in casts
+# 2) Insert an fp16 `flat_weight` if necessary
+def rnn_cast(backend, fn, handle, verbose=False):
+ orig_rnn = utils.get_func(backend, fn)
+ @functools.wraps(orig_rnn)
+ def rnn_wrapper(*args, **kwargs):
+ flat_weight = kwargs.get('flat_weight')
+ if flat_weight is not None:
+ # We replace `flat_weight` with an uninitialized fp16
+ # Tensor. The "actual" weight tensors (provided in `forward`),
+ # will then be set up as ptrs into the buffer and have the
+ # corresponding fp32 values copied in.
+ # We need to call `copy` on the "actual" weights so that the
+ # autograd graph correctly backprops from the wgrads computed
+ # inside cuDNN (on fp16 weights) into the fp32 weights.
+ assert utils.type_string(flat_weight) == 'FloatTensor'
+ if compat.tensor_is_float_tensor() or compat.tensor_is_variable():
+ # Pre-0.4. A little slower, since it zeros out memory.
+ flat_weight_fp16 = flat_weight.new().half().resize_(flat_weight.shape)
+ else:
+ flat_weight_fp16 = torch.empty_like(flat_weight,
+ dtype=torch.float16)
+ kwargs['flat_weight'] = flat_weight_fp16
+ else:
+ flat_weight_fp16 = None
+
+ forward = orig_rnn(*args, **kwargs)
+ @functools.wraps(forward)
+ def fwd_wrapper(*fargs, **fkwargs):
+ assert len(fargs) == 3 or len(fargs) == 4
+ inputs, weights, hiddens = fargs[:3]
+ assert utils.is_fp_tensor(inputs)
+ assert isinstance(weights, list)
+ cast_fn = utils.verbosify(utils.maybe_half,
+ fn,
+ verbose)
+ new_args = []
+
+ # 0) Inputs
+ new_args.append(cast_fn(inputs))
+
+ # 1) Weights
+ if flat_weight_fp16 is not None:
+ fp16_weights = utils.synthesize_flattened_rnn_weights(
+ weights, flat_weight_fp16, fn, verbose)
+ else:
+ fp16_weights = [[cast_fn(w) for w in layer]
+ for layer in weights]
+ new_args.append(fp16_weights)
+
+ # 2) Inputs: either a tuple (for LSTM) or single tensor
+ if isinstance(hiddens, tuple):
+ new_args.append(tuple(cast_fn(x) for x in hiddens))
+ elif utils.is_fp_tensor(hiddens):
+ new_args.append(cast_fn(hiddens))
+ else:
+ # Hiddens can, in principle, be `None` -- pass through
+ new_args.append(hiddens)
+
+ # 3) Batch sizes (0.4 or later only)
+ if len(fargs) == 4:
+ new_args.append(fargs[3])
+
+ return forward(*new_args, **fkwargs)
+ return fwd_wrapper
+ utils.set_func_save(handle, backend, fn, rnn_wrapper)
+
+def new_rnn_cast(fn, handle, verbose=False):
+ # Forward+backward compatibility around https://github.com/pytorch/pytorch/pull/15744
+ # For rnn backend calls that route through _rnn_impls, we must patch the ref
+ # that _rnn_impls stashed. For rnn backend calls that directly invoke
+ # _VF., e.g. _VF.lstm, we can patch onto VariableFunctionsShim,
+ # which in turn has patched the ref named "_VF" in torch.nn.modules.rnn.
+ if utils.has_func(torch.nn.modules.rnn._rnn_impls, fn):
+ mod = torch.nn.modules.rnn._rnn_impls
+ else:
+ mod = torch.nn.modules.rnn._VF
+ assert isinstance(mod, rnn_compat.VariableFunctionsShim)
+ fn = fn.lower()
+ orig_fn = utils.get_func(mod, fn)
+ cast_fn = utils.verbosify(utils.maybe_half, fn, verbose)
+ @functools.wraps(orig_fn)
+ def wrapper(*args, **kwargs):
+ # Exact call signature from modules/rnn.py
+ assert len(args) == 9
+ assert len(kwargs) == 0
+
+ if not _amp_state.handle.is_active():
+ return orig_fn(*args, **kwargs)
+
+ if isinstance(args[6], bool):
+ params_idx = 2 # Not PackedSequence case
+ else:
+ params_idx = 3 # PackedSequence case
+
+ new_args = []
+ for i, arg in enumerate(args):
+ if i == params_idx:
+ num_params = sum([x.numel() for x in arg])
+ fp16_weight_buf = args[0].new_empty((num_params,),
+ dtype=torch.half)
+ casted_weights = utils.new_synthesize_flattened_rnn_weights(
+ arg, fp16_weight_buf, fn, verbose)
+ new_args.append(casted_weights)
+ elif utils.is_fp_tensor(arg):
+ new_args.append(cast_fn(arg))
+ else:
+ new_args.append(arg)
+
+ return orig_fn(*new_args)
+ utils.set_func_save(handle, mod, fn, wrapper)
+
+def disable_casts(mod, fn, handle):
+ if not utils.has_func(mod, fn):
+ return
+
+ orig_fn = utils.get_func(mod, fn)
+ @functools.wraps(orig_fn)
+ def wrapper(*args, **kwargs):
+ with handle._disable_casts():
+ return orig_fn(*args, **kwargs)
+ utils.set_func_save(handle, mod, fn, wrapper)
diff --git a/tools/convert_checkpoint/custom/llavaov_1_5_30b_a3b/__init__.py b/apex/apex/contrib/__init__.py
similarity index 100%
rename from tools/convert_checkpoint/custom/llavaov_1_5_30b_a3b/__init__.py
rename to apex/apex/contrib/__init__.py
diff --git a/apex/apex/contrib/bottleneck/__init__.py b/apex/apex/contrib/bottleneck/__init__.py
new file mode 100644
index 00000000..300b7c36
--- /dev/null
+++ b/apex/apex/contrib/bottleneck/__init__.py
@@ -0,0 +1,2 @@
+from .bottleneck import Bottleneck, SpatialBottleneck
+from .halo_exchangers import HaloExchangerNoComm, HaloExchangerAllGather, HaloExchangerSendRecv, HaloExchangerPeer
diff --git a/apex/apex/contrib/bottleneck/bottleneck.py b/apex/apex/contrib/bottleneck/bottleneck.py
new file mode 100644
index 00000000..5ea5694c
--- /dev/null
+++ b/apex/apex/contrib/bottleneck/bottleneck.py
@@ -0,0 +1,749 @@
+import functools as func
+
+import torch
+import torch.distributed as dist
+from torch import nn
+
+from apex import check_cudnn_version_and_warn
+import fast_bottleneck
+import nccl_p2p_cuda as inc
+
+
+assert check_cudnn_version_and_warn(__name__, 8400)
+
+
+def kaiming_uniform_(tensor, a=0, mode='fan_in', nonlinearity='leaky_relu'):
+ weight_tensor_nchw = tensor
+ nn.init.kaiming_uniform_(weight_tensor_nchw, a=a, mode=mode, nonlinearity=nonlinearity)
+
+def compute_scale_bias_one(nhwc, weight, bias, running_mean, running_var, w_scale, w_bias):
+ scale = weight * running_var.rsqrt()
+ bias = bias - running_mean * scale
+ w_scale.copy_(scale)
+ w_bias.copy_(bias)
+
+def compute_scale_bias_method(nhwc, args):
+ for arg in args:
+ # arg is tuple of (weight, bias, running_mean, running_var, w_scale, w_bias)
+ compute_scale_bias_one(nhwc, *arg)
+
+class FrozenBatchNorm2d(torch.jit.ScriptModule):
+ """
+ BatchNorm2d where the batch statistics and the affine parameters are fixed
+ """
+ def __init__(self, n):
+ super(FrozenBatchNorm2d, self).__init__()
+ self.register_buffer("weight", torch.ones(n))
+ self.register_buffer("bias", torch.zeros(n))
+ self.register_buffer("running_mean", torch.zeros(n))
+ self.register_buffer("running_var", torch.ones(n))
+
+ @torch.jit.script_method
+ def get_scale_bias(self, nhwc):
+ # type: (bool) -> List[torch.Tensor]
+ scale = self.weight * self.running_var.rsqrt()
+ bias = self.bias - self.running_mean * scale
+ if nhwc:
+ scale = scale.reshape(1, 1, 1, -1)
+ bias = bias.reshape(1, 1, 1, -1)
+ else:
+ scale = scale.reshape(1, -1, 1, 1)
+ bias = bias.reshape(1, -1, 1, 1)
+ return scale, bias
+
+ @torch.jit.script_method
+ def forward(self, x):
+ scale, bias = self.get_scale_bias(False)
+ return x * scale + bias
+
+@torch.jit.script
+def drelu_dscale1(grad_o, output, scale1):
+ relu_mask = (output>0)
+ dx_relu = relu_mask * grad_o
+ g1 = dx_relu * scale1
+ return g1, dx_relu
+
+@torch.jit.script
+def drelu_dscale2(grad_o, output, scale1, scale2):
+ relu_mask = (output>0)
+ dx_relu = relu_mask * grad_o
+ g1 = dx_relu * scale1
+ g2 = dx_relu * scale2
+ return g1, g2
+
+class BottleneckFunction(torch.autograd.Function):
+ @staticmethod
+ def forward(ctx, nhwc, stride_1x1, scale, bias, x, *conv):
+ # TODO: clean up order of tensors
+ args = [x, *conv[0:3], *scale[0:3], *bias[0:3]]
+ ctx.downsample = len(conv) > 3
+ if ctx.downsample:
+ args.append(conv[3])
+ args.append(scale[3])
+ args.append(bias[3])
+
+ # weight buffers are always in nhwc while shape can be nhwc or channels_last
+ # here we pass in flag and let c++ handle it
+ # alternatively, we can put all sizes into a fixed format and pass it in
+ outputs = fast_bottleneck.forward(nhwc, stride_1x1, args)
+ ctx.save_for_backward(*(args+outputs))
+ # save relu outputs for drelu
+ ctx.nhwc = nhwc
+ ctx.stride_1x1 = stride_1x1
+ return outputs[2]
+
+ # backward relu is not exposed, MUL with mask used now
+ # only support dgrad
+ @staticmethod
+ def backward(ctx, grad_o):
+ outputs = ctx.saved_tensors[-3:]
+
+ if ctx.downsample:
+ grad_conv3, grad_conv4 = drelu_dscale2(grad_o, outputs[2], ctx.saved_tensors[6], ctx.saved_tensors[11])
+ else:
+ grad_conv3, grad_conv4 = drelu_dscale1(grad_o, outputs[2], ctx.saved_tensors[6])
+
+ # create input vector for backward
+ t_list = [*ctx.saved_tensors[0:10]]
+ t_list.append(grad_conv3)
+ t_list.append(grad_conv4)
+
+ # outputs used for wgrad and generating drelu mask
+ t_list.append(outputs[0])
+ t_list.append(outputs[1])
+
+ # in case there is downsample
+ if ctx.downsample:
+ t_list.append(ctx.saved_tensors[10])
+
+ grads = fast_bottleneck.backward(ctx.nhwc, ctx.stride_1x1, t_list)
+
+ return (None, None, None, None, *grads)
+
+bottleneck_function = BottleneckFunction.apply
+
+def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
+ """3x3 convolution with padding"""
+ return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+ padding=dilation, groups=groups, bias=False, dilation=dilation)
+
+def conv1x1(in_planes, out_planes, stride=1):
+ """1x1 convolution"""
+ return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
+
+class Bottleneck(torch.nn.Module):
+ # Bottleneck in torchvision places the stride for downsampling at 3x3 convolution(self.conv2)
+ # while original implementation places the stride at the first 1x1 convolution(self.conv1)
+ # according to "Deep residual learning for image recognition"https://arxiv.org/abs/1512.03385.
+ # This variant is also known as ResNet V1.5 and improves accuracy according to
+ # https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch.
+ # here we put it at 1x1
+
+ def __init__(self, in_channels, bottleneck_channels, out_channels, stride=1, groups=1,
+ dilation=1, norm_func=None, use_cudnn=False, explicit_nhwc=False):
+ super(Bottleneck, self).__init__()
+ if groups != 1:
+ raise RuntimeError('Only support groups == 1')
+ if dilation != 1:
+ raise RuntimeError('Only support dilation == 1')
+ if norm_func == None:
+ norm_func = FrozenBatchNorm2d
+ else:
+ raise RuntimeError('Only support frozen BN now.')
+
+ if stride != 1 or in_channels != out_channels:
+ self.downsample = nn.Sequential(
+ conv1x1(in_channels, out_channels, stride),
+ norm_func(out_channels),
+ )
+ else:
+ self.downsample = None
+
+ # Both self.conv2 and self.downsample layers downsample the input when stride != 1
+ self.conv1 = conv1x1(in_channels, bottleneck_channels, stride)
+ self.conv2 = conv3x3(bottleneck_channels, bottleneck_channels)
+ self.conv3 = conv1x1(bottleneck_channels, out_channels)
+ self.relu = nn.ReLU(inplace=True)
+ self.stride = stride
+
+ self.bn1 = norm_func(bottleneck_channels)
+ self.bn2 = norm_func(bottleneck_channels)
+ self.bn3 = norm_func(out_channels)
+ self.w_scale = None
+
+ self.use_cudnn = use_cudnn
+
+ # setup conv weights
+ self.w_conv = [self.conv1.weight, self.conv2.weight, self.conv3.weight]
+ if self.downsample is not None:
+ self.w_conv.append(self.downsample[0].weight)
+
+ # init weight in nchw format before possible transpose
+ for w in self.w_conv:
+ kaiming_uniform_(w, a=1)
+
+ # TODO: prevent unsupported case usage
+ # support cases
+ # native cudnn
+ # normal yes no
+ # channel_last yes yes
+ # explicit_nhwc no yes
+ self.explicit_nhwc = explicit_nhwc
+ if self.explicit_nhwc:
+ for p in self.parameters():
+ with torch.no_grad():
+ p.data = p.data.permute(0,2,3,1).contiguous()
+
+ return
+
+ # Returns single callable that recomputes scale and bias for all frozen batch-norms.
+ # This method must be called before cuda graphing.
+ # The callable it returns can be called anytime.
+ # Calling this method will prevent these from being computed every forward call.
+ def get_scale_bias_callable(self):
+ self.w_scale, self.w_bias, args = [], [], []
+ batch_norms = [self.bn1, self.bn2, self.bn3]
+ if self.downsample is not None:
+ batch_norms.append(self.downsample[1])
+ for bn in batch_norms:
+ s = torch.empty_like(bn.weight)
+ b = torch.empty_like(s)
+ args.append( (bn.weight, bn.bias, bn.running_mean, bn.running_var, s, b) )
+ if self.explicit_nhwc:
+ self.w_scale.append( s.reshape(1, 1, 1, -1) )
+ self.w_bias.append( b.reshape(1, 1, 1, -1) )
+ else:
+ self.w_scale.append( s.reshape(1, -1, 1, 1) )
+ self.w_bias.append( b.reshape(1, -1, 1, 1) )
+ return func.partial(compute_scale_bias_method, self.explicit_nhwc, args)
+
+ def forward(self, x):
+ if self.use_cudnn:
+ if self.w_scale is None:
+ # calculate scale/bias from registered buffers
+ # TODO: make this better
+ s1, b1 = self.bn1.get_scale_bias(self.explicit_nhwc)
+ s2, b2 = self.bn2.get_scale_bias(self.explicit_nhwc)
+ s3, b3 = self.bn3.get_scale_bias(self.explicit_nhwc)
+ w_scale = [s1, s2, s3]
+ w_bias = [b1, b2, b3]
+ if self.downsample is not None:
+ s4, b4 = self.downsample[1].get_scale_bias(self.explicit_nhwc)
+ w_scale.append(s4)
+ w_bias.append(b4)
+ out = bottleneck_function(self.explicit_nhwc, self.stride, w_scale, w_bias, x, *self.w_conv)
+ else:
+ out = bottleneck_function(self.explicit_nhwc, self.stride, self.w_scale, self.w_bias, x, *self.w_conv)
+ return out
+
+ if self.explicit_nhwc:
+ raise RuntimeError('explicit nhwc with native ops is not supported.')
+
+ # fallback to native ops
+ identity = x
+
+ out = self.conv1(x)
+ out = self.bn1(out)
+ out = self.relu(out)
+
+ out = self.conv2(out)
+ out = self.bn2(out)
+ out = self.relu(out)
+
+ out = self.conv3(out)
+ out = self.bn3(out)
+
+ if self.downsample is not None:
+ identity = self.downsample(x)
+
+ out += identity
+ out = self.relu(out)
+
+ return out
+
+
+class SpatialBottleneckFunction(torch.autograd.Function):
+ @staticmethod
+ def forward(ctx, spatial_group_size, spatial_group_rank, spatial_communicator, spatial_halo_exchanger, spatial_method, use_delay_kernel, explicit_nhwc, stride_1x1, scale, bias, thresholdTop, thresholdBottom, x, *conv):
+ if spatial_group_size > 1:
+ stream1 = spatial_halo_exchanger.stream1
+ stream2 = spatial_halo_exchanger.stream2
+ stream3 = spatial_halo_exchanger.stream3
+
+ # TODO: clean up order of tensors
+ args = [x, *conv[0:3], *scale[0:3], *bias[0:3]]
+ ctx.downsample = len(conv) > 3
+ if ctx.downsample:
+ args.append(conv[3])
+ args.append(scale[3])
+ args.append(bias[3])
+
+ # weight buffers are always in explicit_nhwc while shape can be explicit_nhwc or channels_last
+ # here we pass in flag and let c++ handle it
+ # alternatively, we can put all sizes into a fixed format and pass it in
+ outputs = fast_bottleneck.forward_init(explicit_nhwc, stride_1x1, args)
+ fast_bottleneck.forward_out1(explicit_nhwc, stride_1x1, args, outputs)
+
+ if spatial_group_size > 1:
+ out1 = outputs[0]
+ if explicit_nhwc:
+ N,Hs,W,C = list(out1.shape)
+ memory_format = torch.contiguous_format
+ out1_pad = torch.empty([N,Hs+2,W,C], dtype=out1.dtype, device='cuda')
+ else:
+ N,C,Hs,W = list(out1.shape)
+ memory_format = torch.channels_last if out1.is_contiguous(memory_format=torch.channels_last) else torch.contiguous_format
+ out1_pad = torch.empty([N,C,Hs+2,W], dtype=out1.dtype, device='cuda', memory_format=memory_format)
+ stream1.wait_stream(torch.cuda.current_stream())
+ if spatial_method != 2: stream3.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(stream1):
+ if explicit_nhwc:
+ top_out1_halo = out1_pad[:,:1,:,:]
+ btm_out1_halo = out1_pad[:,Hs+1:Hs+2,:,:]
+ spatial_halo_exchanger.left_right_halo_exchange(out1[:,:1,:,:], out1[:,Hs-1:,:,:], top_out1_halo, btm_out1_halo)
+ else:
+ top_out1_halo = out1_pad[:,:,:1,:]
+ btm_out1_halo = out1_pad[:,:,Hs+1:Hs+2,:]
+ spatial_halo_exchanger.left_right_halo_exchange(out1[:,:,:1,:], out1[:,:,Hs-1:,:], top_out1_halo, btm_out1_halo)
+ if spatial_method == 1:
+ # overlap mid convolution with halo transfer
+ if spatial_group_rank < spatial_group_size-1:
+ stream2.wait_stream(stream1)
+ with torch.cuda.stream(stream2):
+ if explicit_nhwc:
+ btm_fat_halo = torch.empty((N,3,W,C),dtype=out1.dtype,device=out1.device)
+ btm_fat_halo[:,0:2,:,:].copy_(out1[:,Hs-2:,:,:])
+ btm_fat_halo[:,2:,:,:].copy_(btm_out1_halo)
+ else:
+ btm_fat_halo = torch.empty((N,C,3,W),dtype=out1.dtype,device=out1.device)
+ btm_fat_halo[:,:,0:2,:].copy_(out1[:,:,Hs-2:,:])
+ btm_fat_halo[:,:,2:,:].copy_(btm_out1_halo)
+ btm_out2 = fast_bottleneck.forward_out2_halo(explicit_nhwc, btm_fat_halo, args)
+ if spatial_group_rank > 0:
+ with torch.cuda.stream(stream1):
+ if explicit_nhwc:
+ top_fat_halo = torch.empty((N,3,W,C),dtype=out1.dtype,device=out1.device)
+ top_fat_halo[:,:1,:,:].copy_(top_out1_halo)
+ top_fat_halo[:,1:3,:,:].copy_(out1[:,:2,:,:])
+ else:
+ top_fat_halo = torch.empty((N,C,3,W),dtype=out1.dtype,device=out1.device)
+ top_fat_halo[:,:,:1,:].copy_(top_out1_halo)
+ top_fat_halo[:,:,1:3,:].copy_(out1[:,:,:2,:])
+ top_out2 = fast_bottleneck.forward_out2_halo(explicit_nhwc, top_fat_halo, args)
+ if use_delay_kernel: inc.add_delay(10)
+ elif spatial_method != 2 and spatial_method != 3:
+ assert(False), "spatial_method must be 1, 2 or 3"
+
+ if spatial_group_size <= 1:
+ fast_bottleneck.forward_out2(explicit_nhwc, stride_1x1, args, outputs)
+ elif spatial_method == 1:
+ fast_bottleneck.forward_out2(explicit_nhwc, stride_1x1, args, outputs)
+ with torch.cuda.stream(stream3):
+ if explicit_nhwc:
+ out1_pad[:,1:Hs+1,:,:].copy_(out1)
+ else:
+ out1_pad[:,:,1:Hs+1,:].copy_(out1)
+ elif spatial_method == 2:
+ # wait for halo transfer to finish before doing a full convolution of padded x
+ if explicit_nhwc:
+ out1_pad[:,1:Hs+1,:,:].copy_(out1)
+ else:
+ out1_pad[:,:,1:Hs+1,:].copy_(out1)
+ torch.cuda.current_stream().wait_stream(stream1)
+ fast_bottleneck.forward_out2_pad(explicit_nhwc, stride_1x1, args, outputs, out1_pad)
+ elif spatial_method == 3:
+ fast_bottleneck.forward_out2_mask(explicit_nhwc, stride_1x1, args, outputs, thresholdTop, thresholdBottom)
+ with torch.cuda.stream(stream3):
+ if explicit_nhwc:
+ out1_pad[:,1:Hs+1,:,:].copy_(out1)
+ else:
+ out1_pad[:,:,1:Hs+1,:].copy_(out1)
+
+ # compute halo cells for outputs[1] (out2)
+ if spatial_group_size > 1:
+ out2 = outputs[1]
+ if explicit_nhwc:
+ top_out2_halo = out2[:,:1,:,:]
+ btm_out2_halo = out2[:,Hs-1:,:,:]
+ else:
+ top_out2_halo = out2[:,:,:1,:]
+ btm_out2_halo = out2[:,:,Hs-1:,:]
+ if spatial_method == 1:
+ if spatial_group_rank > 0:
+ torch.cuda.current_stream().wait_stream(stream1)
+ top_out2_halo.copy_(top_out2)
+ if spatial_group_rank < spatial_group_size-1:
+ torch.cuda.current_stream().wait_stream(stream2)
+ btm_out2_halo.copy_(btm_out2)
+ elif spatial_method == 3:
+ # Note
+ # out2 halo correction cannot overlap with anything since it has
+ # to wait for out2_mask to finish, but itself has to finish before
+ # the first kernel of _forward_rest can launch.
+ # At least we can overlap the two halo correction kernels.
+ if spatial_group_rank < spatial_group_size-1:
+ stream2.wait_stream(stream1) # wait for halo transfers to finish
+ stream2.wait_stream(torch.cuda.current_stream()) # wait for *_out2_mask to finish
+ with torch.cuda.stream(stream2):
+ w1by3 = args[2][:,2:3,:,:].clone()
+ btm_out1_halo = btm_out1_halo.clone()
+ btm_out2 = fast_bottleneck.forward_out2_halo_corr(explicit_nhwc, btm_out1_halo, args, w1by3, btm_out2_halo.clone())
+ btm_out2_halo.copy_(btm_out2)
+ if spatial_group_rank > 0:
+ stream1.wait_stream(torch.cuda.current_stream()) # wait for *_out2_mask to finish
+ with torch.cuda.stream(stream1):
+ w1by3 = args[2][:,:1,:,:].clone()
+ top_out1_halo = top_out1_halo.clone()
+ top_out2 = fast_bottleneck.forward_out2_halo_corr(explicit_nhwc, top_out1_halo, args, w1by3, top_out2_halo.clone())
+ top_out2_halo.copy_(top_out2)
+ if spatial_group_rank < spatial_group_size-1:
+ torch.cuda.current_stream().wait_stream(stream2)
+ if spatial_group_rank > 0:
+ torch.cuda.current_stream().wait_stream(stream1)
+
+ fast_bottleneck.forward_rest(explicit_nhwc, stride_1x1, args, outputs)
+ # save halos for backward pass
+ if spatial_group_size > 1:
+ if spatial_method != 2:
+ # make sure copy of mid-section of out1 into out1_pad is done before exiting
+ torch.cuda.current_stream().wait_stream(stream3)
+ ctx.save_for_backward(*(args+outputs+[out1_pad,]))
+ else:
+ ctx.save_for_backward(*(args+outputs))
+ # save relu outputs for drelu
+ ctx.explicit_nhwc = explicit_nhwc
+ ctx.stride_1x1 = stride_1x1
+ ctx.spatial_group_size = spatial_group_size
+ if spatial_group_size > 1:
+ ctx.spatial_group_rank = spatial_group_rank
+ ctx.spatial_halo_exchanger = spatial_halo_exchanger
+ ctx.spatial_method = spatial_method
+ ctx.use_delay_kernel = use_delay_kernel
+ ctx.thresholdTop = thresholdTop
+ ctx.thresholdBottom = thresholdBottom
+ ctx.stream1 = stream1
+ ctx.stream2 = stream2
+ ctx.stream3 = stream3
+ return outputs[2]
+
+ # backward relu is not exposed, MUL with mask used now
+ # only support dgrad
+ @staticmethod
+ def backward(ctx, grad_o):
+ if ctx.spatial_group_size > 1:
+ out1_pad = ctx.saved_tensors[-1]
+ outputs = ctx.saved_tensors[-4:-1]
+ else:
+ outputs = ctx.saved_tensors[-3:]
+
+ if ctx.downsample:
+ grad_conv3, grad_conv4 = drelu_dscale2(grad_o, outputs[2], ctx.saved_tensors[6], ctx.saved_tensors[11])
+ else:
+ grad_conv3, grad_conv4 = drelu_dscale1(grad_o, outputs[2], ctx.saved_tensors[6])
+
+ # create input vector for backward
+ t_list = [*ctx.saved_tensors[0:10]]
+ t_list.append(grad_conv3)
+ t_list.append(grad_conv4)
+
+ # outputs used for wgrad and generating drelu mask
+ t_list.append(outputs[0])
+ t_list.append(outputs[1])
+
+ # in case there is downsample
+ if ctx.downsample:
+ t_list.append(ctx.saved_tensors[10])
+
+ grads = fast_bottleneck.backward_init(ctx.explicit_nhwc, ctx.stride_1x1, t_list)
+ wgrad3_stream = torch.cuda.Stream()
+ wgrad3_stream.wait_stream(torch.cuda.current_stream())
+ grad_out2 = fast_bottleneck.backward_grad_out2(ctx.explicit_nhwc, ctx.stride_1x1, t_list, grads)
+ wgrad2_stream = torch.cuda.Stream()
+ wgrad2_stream.wait_stream(torch.cuda.current_stream())
+ # do halo exchange of grad_out2 here
+ # compute halo cells for grad_out1
+ if ctx.spatial_group_size > 1:
+ if ctx.explicit_nhwc:
+ N,Hs,W,C = list(grad_out2.shape)
+ else:
+ N,C,Hs,W = list(grad_out2.shape)
+ relu1 = t_list[12]
+ ctx.stream1.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(ctx.stream1):
+ top_halo, btm_halo = ctx.spatial_halo_exchanger.left_right_halo_exchange(grad_out2[:,:1,:,:], grad_out2[:,Hs-1:,:,:])
+ # copy halos to send buffer
+ if ctx.spatial_method == 1 or ctx.spatial_method == 2:
+ # 1 -> halo recompute approach
+ # 2 -> wait for concatenated halos, then do single conv on full input (not implemented yet for bprop)
+ if ctx.spatial_group_rank < ctx.spatial_group_size-1:
+ ctx.stream2.wait_stream(ctx.stream1)
+ with torch.cuda.stream(ctx.stream2):
+ if ctx.explicit_nhwc:
+ btm_fat_halo = torch.empty((N,3,W,C),dtype=grad_out2.dtype,device=grad_out2.device)
+ btm_fat_halo[:,:2,:,:].copy_(grad_out2[:,Hs-2:,:,:])
+ btm_fat_halo[:,2:,:,:].copy_(btm_halo)
+ btm_fat_relu_halo = torch.empty((N,3,W,C),dtype=grad_out2.dtype,device=grad_out2.device)
+ btm_fat_relu_halo[:,:2,:,:].copy_(relu1[:,Hs-2:,:,:])
+ btm_fat_relu_halo[:,2:,:,:].zero_()
+ else:
+ btm_fat_halo = torch.empty((N,C,3,W),dtype=grad_out2.dtype,device=grad_out2.device)
+ btm_fat_halo[:,:,:2,:].copy_(grad_out2[:,:,Hs-2:,:])
+ btm_fat_halo[:,:,2:,:].copy_(btm_halo)
+ btm_fat_relu_halo = torch.empty((N,C,3,W),dtype=grad_out2.dtype,device=grad_out2.device)
+ btm_fat_relu_halo[:,:,:2,:].copy_(relu1[:,:,Hs-2:,:])
+ btm_fat_relu_halo[:,:,2:,:].zero_()
+ btm_grad_out1_halo = fast_bottleneck.backward_grad_out1_halo(ctx.explicit_nhwc, ctx.stride_1x1, t_list, grads, btm_fat_halo, btm_fat_relu_halo)
+ if ctx.explicit_nhwc:
+ btm_grad_out1_halo = btm_grad_out1_halo[:,1:2,:,:]
+ else:
+ btm_grad_out1_halo = btm_grad_out1_halo[:,:,1:2,:]
+ if ctx.spatial_group_rank > 0:
+ with torch.cuda.stream(ctx.stream1):
+ if ctx.explicit_nhwc:
+ top_fat_halo = torch.empty((N,3,W,C),dtype=grad_out2.dtype,device=grad_out2.device)
+ top_fat_halo[:,:1,:,:].copy_(top_halo)
+ top_fat_halo[:,1:,:,:].copy_(grad_out2[:,:2,:,:])
+ top_fat_relu_halo = torch.empty((N,3,W,C),dtype=grad_out2.dtype,device=grad_out2.device)
+ top_fat_relu_halo[:,:1,:,:].zero_()
+ top_fat_relu_halo[:,1:,:,:].copy_(relu1[:,:2,:,:])
+ else:
+ top_fat_halo = torch.empty((N,C,3,W),dtype=grad_out2.dtype,device=grad_out2.device)
+ top_fat_halo[:,:,:1,:].copy_(top_halo)
+ top_fat_halo[:,:,1:,:].copy_(grad_out2[:,:,:2,:])
+ top_fat_relu_halo = torch.empty((N,C,3,W),dtype=grad_out2.dtype,device=grad_out2.device)
+ top_fat_relu_halo[:,:,:1,:].zero_()
+ top_fat_relu_halo[:,:,1:,:].copy_(relu1[:,:,:2,:])
+ top_grad_out1_halo = fast_bottleneck.backward_grad_out1_halo(ctx.explicit_nhwc, ctx.stride_1x1, t_list, grads, top_fat_halo, top_fat_relu_halo)
+ if ctx.explicit_nhwc:
+ top_grad_out1_halo = top_grad_out1_halo[:,1:2,:,:]
+ else:
+ top_grad_out1_halo = top_grad_out1_halo[:,:,1:2,:]
+ if ctx.use_delay_kernel: inc.add_delay(10)
+ elif ctx.spatial_method != 3:
+ assert(False), "spatial_method must be 1, 2 or 3"
+
+ # compute grad_out1 for internal cells
+ if ctx.spatial_group_size <= 1 or ctx.spatial_method == 1 or ctx.spatial_method == 2:
+ grad_out1 = fast_bottleneck.backward_grad_out1(ctx.explicit_nhwc, ctx.stride_1x1, t_list, grads, grad_out2)
+ elif ctx.spatial_group_size > 1 and ctx.spatial_method == 3:
+ grad_out1 = fast_bottleneck.backward_grad_out1_mask(ctx.explicit_nhwc, ctx.stride_1x1, t_list, grads, grad_out2, ctx.thresholdTop, ctx.thresholdBottom)
+
+ # apply halo cells to grad_out1
+ if ctx.spatial_group_size > 1:
+ w = t_list[2]
+ z = t_list[4]
+ relu1 = t_list[12]
+ #print("w.shape = %s, z.shape = %s, relu1.shape = %s" % (str(list(w.shape)), str(list(z.shape)), str(list(relu1.shape))))
+ if ctx.spatial_method == 1 or ctx.spatial_method == 2:
+ if ctx.spatial_group_rank < ctx.spatial_group_size-1:
+ torch.cuda.current_stream().wait_stream(ctx.stream2)
+ if ctx.explicit_nhwc:
+ grad_out1[:,Hs-1:,:,:].copy_(btm_grad_out1_halo)
+ else:
+ grad_out1[:,:,Hs-1:,:].copy_(btm_grad_out1_halo)
+ #print("ctx.spatial_group_rank = %d, apply grad_out1 btm halo (grad_out1.shape = %s)" % (ctx.spatial_group_rank, str(list(grad_out1.shape))))
+ if ctx.spatial_group_rank > 0:
+ torch.cuda.current_stream().wait_stream(ctx.stream1)
+ if ctx.explicit_nhwc:
+ grad_out1[:,:1,:,:].copy_(top_grad_out1_halo)
+ else:
+ grad_out1[:,:,:1,:].copy_(top_grad_out1_halo)
+ #print("ctx.spatial_group_rank = %d, apply grad_out1 top halo (grad_out1.shape = %s)" % (ctx.spatial_group_rank, str(list(grad_out1.shape))))
+ elif ctx.spatial_method == 3:
+ if ctx.spatial_group_rank < ctx.spatial_group_size-1:
+ if ctx.explicit_nhwc:
+ btm_relu_halo = relu1[:,Hs-1:,:,:].clone()
+ btm_grad_out1 = grad_out1[:,Hs-1:,:,:]
+ else:
+ btm_relu_halo = relu1[:,:,Hs-1:,:].clone()
+ btm_grad_out1 = grad_out1[:,:,Hs-1:,:]
+ w1by3 = w[:,:1,:,:].clone()
+ ctx.stream2.wait_stream(ctx.stream1) # wait for halo transfers to finish
+ ctx.stream2.wait_stream(torch.cuda.current_stream()) # wait for backward_grad_out1_mask to finish before launching halo correction kernel
+ with torch.cuda.stream(ctx.stream2):
+ btm_grad_out1_halo = fast_bottleneck.backward_grad_out1_halo_corr(ctx.explicit_nhwc, ctx.stride_1x1, t_list, w1by3, grads, btm_halo, btm_relu_halo, btm_grad_out1.clone())
+ btm_grad_out1.copy_(btm_grad_out1_halo)
+ if ctx.spatial_group_rank > 0:
+ if ctx.explicit_nhwc:
+ top_relu_halo = relu1[:,:1,:,:].clone()
+ top_grad_out1 = grad_out1[:,:1,:,:]
+ else:
+ top_relu_halo = relu1[:,:,:1,:].clone()
+ top_grad_out1 = grad_out1[:,:,:1,:]
+ w1by3 = w[:,2:,:,:].clone()
+ ctx.stream1.wait_stream(torch.cuda.current_stream()) # wait for backward_grad_out1_mask to finish before launching halo correction kernel
+ with torch.cuda.stream(ctx.stream1):
+ top_grad_out1_halo = fast_bottleneck.backward_grad_out1_halo_corr(ctx.explicit_nhwc, ctx.stride_1x1, t_list, w1by3, grads, top_halo, top_relu_halo, top_grad_out1.clone())
+ top_grad_out1.copy_(top_grad_out1_halo)
+ if ctx.spatial_group_rank < ctx.spatial_group_size-1:
+ torch.cuda.current_stream().wait_stream(ctx.stream2) # wait for halo correction to finish
+ if ctx.spatial_group_rank > 0:
+ torch.cuda.current_stream().wait_stream(ctx.stream1)
+
+ wgrad1_stream = torch.cuda.Stream()
+ wgrad1_stream.wait_stream(torch.cuda.current_stream())
+ fast_bottleneck.backward_rest(ctx.explicit_nhwc, ctx.stride_1x1, t_list, grads, grad_out2, grad_out1)
+ with torch.cuda.stream(wgrad3_stream):
+ fast_bottleneck.backward_wgrad3(ctx.explicit_nhwc, ctx.stride_1x1, t_list, grads)
+ with torch.cuda.stream(wgrad2_stream):
+ if ctx.spatial_group_size > 1:
+ fast_bottleneck.backward_wgrad2_pad(ctx.explicit_nhwc, ctx.stride_1x1, t_list, grads, out1_pad, grad_out2)
+ else:
+ fast_bottleneck.backward_wgrad2(ctx.explicit_nhwc, ctx.stride_1x1, t_list, grads, grad_out2)
+ with torch.cuda.stream(wgrad1_stream):
+ fast_bottleneck.backward_wgrad1(ctx.explicit_nhwc, ctx.stride_1x1, t_list, grads, grad_out1)
+ torch.cuda.current_stream().wait_stream(wgrad3_stream)
+ torch.cuda.current_stream().wait_stream(wgrad2_stream)
+ torch.cuda.current_stream().wait_stream(wgrad1_stream)
+
+ return (None, None, None, None, None, None, None, None, None, None, None, None, *grads)
+
+spatial_bottleneck_function = SpatialBottleneckFunction.apply
+
+class SpatialBottleneck(torch.nn.Module):
+ # Bottleneck in torchvision places the stride for downsampling at 3x3 convolution(self.conv2)
+ # while original implementation places the stride at the first 1x1 convolution(self.conv1)
+ # according to "Deep residual learning for image recognition"https://arxiv.org/abs/1512.03385.
+ # This variant is also known as ResNet V1.5 and improves accuracy according to
+ # https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch.
+ # here we put it at 1x1
+
+ def __init__(self, in_channels, bottleneck_channels, out_channels, stride=1, groups=1,
+ dilation=1, norm_func=None, use_cudnn=False, explicit_nhwc=False,
+ spatial_parallel_args=None):
+ super(SpatialBottleneck, self).__init__()
+ if groups != 1:
+ raise RuntimeError('Only support groups == 1')
+ if dilation != 1:
+ raise RuntimeError('Only support dilation == 1')
+ if norm_func == None:
+ norm_func = FrozenBatchNorm2d
+ else:
+ raise RuntimeError('Only support frozen BN now.')
+
+ if stride != 1 or in_channels != out_channels:
+ self.downsample = nn.Sequential(
+ conv1x1(in_channels, out_channels, stride),
+ norm_func(out_channels),
+ )
+ else:
+ self.downsample = None
+
+ # Both self.conv2 and self.downsample layers downsample the input when stride != 1
+ self.conv1 = conv1x1(in_channels, bottleneck_channels, stride)
+ self.conv2 = conv3x3(bottleneck_channels, bottleneck_channels)
+ self.conv3 = conv1x1(bottleneck_channels, out_channels)
+ self.relu = nn.ReLU(inplace=True)
+ self.stride = stride
+
+ self.bn1 = norm_func(bottleneck_channels)
+ self.bn2 = norm_func(bottleneck_channels)
+ self.bn3 = norm_func(out_channels)
+ self.w_scale = None
+
+ self.use_cudnn = use_cudnn
+
+ # setup conv weights
+ self.w_conv = [self.conv1.weight, self.conv2.weight, self.conv3.weight]
+ if self.downsample is not None:
+ self.w_conv.append(self.downsample[0].weight)
+
+ # init weight in nchw format before possible transpose
+ for w in self.w_conv:
+ kaiming_uniform_(w, a=1)
+
+ self.thresholdTop, self.thresholdBottom = None, None
+
+ # TODO: prevent unsupported case usage
+ # support cases
+ # native cudnn
+ # normal yes no
+ # channel_last yes yes
+ # explicit_nhwc no yes
+ self.explicit_nhwc = explicit_nhwc
+ if self.explicit_nhwc:
+ for p in self.parameters():
+ with torch.no_grad():
+ p.data = p.data.permute(0,2,3,1).contiguous()
+
+ # spatial communicator
+ if spatial_parallel_args is None:
+ self.spatial_parallel_args = (1, 0, None, None, 0, False)
+ else:
+ self.spatial_parallel_args = spatial_parallel_args
+ return
+
+ # Returns single callable that recomputes scale and bias for all frozen batch-norms.
+ # This method must be called before cuda graphing.
+ # The callable it returns can be called anytime.
+ # Calling this method will prevent these from being computed every forward call.
+ def get_scale_bias_callable(self):
+ self.w_scale, self.w_bias, args = [], [], []
+ batch_norms = [self.bn1, self.bn2, self.bn3]
+ if self.downsample is not None:
+ batch_norms.append(self.downsample[1])
+ for bn in batch_norms:
+ s = torch.empty_like(bn.weight)
+ b = torch.empty_like(s)
+ args.append( (bn.weight, bn.bias, bn.running_mean, bn.running_var, s, b) )
+ if self.explicit_nhwc:
+ self.w_scale.append( s.reshape(1, 1, 1, -1) )
+ self.w_bias.append( b.reshape(1, 1, 1, -1) )
+ else:
+ self.w_scale.append( s.reshape(1, -1, 1, 1) )
+ self.w_bias.append( b.reshape(1, -1, 1, 1) )
+ return func.partial(compute_scale_bias_method, self.explicit_nhwc, args)
+
+ def forward(self, x):
+ if self.use_cudnn:
+ if self.thresholdTop is None:
+ spatial_group_size, spatial_group_rank, _, _, _, _ = self.spatial_parallel_args
+ if self.explicit_nhwc:
+ N,H,W,C = list(x.shape)
+ else:
+ N,C,H,W = list(x.shape)
+ self.thresholdTop = torch.tensor([1 if spatial_group_rank > 0 else 0], dtype=torch.int32, device='cuda')
+ self.thresholdBottom = torch.tensor([H-2 if spatial_group_rank < spatial_group_size - 1 else H-1], dtype=torch.int32, device='cuda')
+
+ if self.w_scale is None:
+ # calculate scale/bias from registered buffers
+ # TODO: make this better
+ s1, b1 = self.bn1.get_scale_bias(self.explicit_nhwc)
+ s2, b2 = self.bn2.get_scale_bias(self.explicit_nhwc)
+ s3, b3 = self.bn3.get_scale_bias(self.explicit_nhwc)
+ w_scale = [s1, s2, s3]
+ w_bias = [b1, b2, b3]
+ if self.downsample is not None:
+ s4, b4 = self.downsample[1].get_scale_bias(self.explicit_nhwc)
+ w_scale.append(s4)
+ w_bias.append(b4)
+ out = spatial_bottleneck_function(*self.spatial_parallel_args, self.explicit_nhwc, self.stride, w_scale, w_bias, self.thresholdTop, self.thresholdBottom, x, *self.w_conv)
+ else:
+ out = spatial_bottleneck_function(*self.spatial_parallel_args, self.explicit_nhwc, self.stride, self.w_scale, self.w_bias, self.thresholdTop, self.thresholdBottom, x, *self.w_conv)
+ return out
+
+ if self.explicit_nhwc:
+ raise RuntimeError('explicit nhwc with native ops is not supported.')
+
+ # fallback to native ops
+ identity = x
+
+ out = self.conv1(x)
+ out = self.bn1(out)
+ out = self.relu(out)
+
+ out = self.conv2(out)
+ out = self.bn2(out)
+ out = self.relu(out)
+
+ out = self.conv3(out)
+ out = self.bn3(out)
+
+ if self.downsample is not None:
+ identity = self.downsample(x)
+
+ out += identity
+ out = self.relu(out)
+
+ return out
+
diff --git a/apex/apex/contrib/bottleneck/halo_exchangers.py b/apex/apex/contrib/bottleneck/halo_exchangers.py
new file mode 100644
index 00000000..3299e823
--- /dev/null
+++ b/apex/apex/contrib/bottleneck/halo_exchangers.py
@@ -0,0 +1,180 @@
+import torch
+import torch.distributed as dist
+from torch import nn
+import nccl_p2p_cuda as inc
+import peer_memory_cuda as pm
+
+# Communication free halo exchanger.
+# NB! This halo exchanger does not exchange halos with neighbors as it should, it merely swaps the inputs
+# NB! This is only useful for performance testing.
+# NB! Do not use for actual production runs
+class HaloExchanger(object):
+ def __init__(self, ranks, rank_in_group):
+ self.stream1 = torch.cuda.Stream()
+ self.stream2 = torch.cuda.Stream()
+ self.stream3 = torch.cuda.Stream()
+ self.group_size = len(ranks)
+ self.ranks = ranks
+ self.rank_in_group = rank_in_group
+ self.wrap_around_left_rank_in_group = (rank_in_group + self.group_size - 1) % self.group_size
+ self.wrap_around_right_rank_in_group = (rank_in_group + 1) % self.group_size
+ self.left_rank = ranks[rank_in_group-1] if rank_in_group > 0 else -1
+ self.left_zero = True if rank_in_group == 0 else False
+ self.right_rank = ranks[rank_in_group+1] if rank_in_group < self.group_size - 1 else -1
+ self.right_zero = True if rank_in_group == self.group_size - 1 else False
+
+class HaloExchangerNoComm(HaloExchanger):
+ def __init__(self, ranks, rank_in_group):
+ super(HaloExchangerNoComm, self).__init__(ranks, rank_in_group)
+
+ def left_right_halo_exchange(self, left_output_halo, right_output_halo, left_input_halo=None, right_input_halo=None):
+ if left_input_halo is None:
+ return right_output_halo, left_output_halo
+ else:
+ left_input_halo.copy_(right_output_halo)
+ right_input_halo.copy_(left_output_halo)
+
+class HaloExchangerAllGather(HaloExchanger):
+ def __init__(self, ranks, rank_in_group, comm):
+ super(HaloExchangerAllGather, self).__init__(ranks, rank_in_group)
+ # self.comm must be NCCL process_group created with torch.distributed.new_group(ranks=ranks)
+ self.comm = comm
+
+ def left_right_halo_exchange(self, left_output_halo, right_output_halo, left_input_halo=None, right_input_halo=None):
+ N,Hh,W,C = list(left_output_halo.shape)
+ send_halos = torch.empty((N,2*Hh,W,C),dtype=left_output_halo.dtype,device=left_output_halo.device)
+ send_halos[:,:Hh,:,:].copy_(left_output_halo)
+ send_halos[:,Hh:,:,:].copy_(right_output_halo)
+ all_halos = torch.empty((N,2*Hh*self.group_size,W,C),dtype=left_output_halo.dtype,device=left_output_halo.device)
+ all_halos = [all_halos[:,i*2*Hh:(i+1)*2*Hh,:,:] for i in range(self.group_size)]
+ torch.distributed.all_gather(all_halos,send_halos,group=self.comm,no_copy=True)
+ ag_left_input_halo = all_halos[self.wrap_around_left_rank_in_group][:,Hh:,:,:]
+ ag_right_input_halo = all_halos[self.wrap_around_right_rank_in_group][:,:Hh,:,:]
+ if left_input_halo is None:
+ if self.left_zero:
+ ag_left_input_halo.zero_()
+ if self.right_zero:
+ ag_right_input_halo.zero_()
+ return ag_left_input_halo, ag_right_input_halo
+ else:
+ if self.left_zero:
+ left_input_halo.zero_()
+ else:
+ left_input_halo.copy_(ag_left_input_halo)
+ if self.right_zero:
+ right_input_halo.zero_()
+ else:
+ right_input_halo.copy_(ag_right_input_halo)
+
+class HaloExchangerSendRecv(HaloExchanger):
+ def __init__(self, ranks, rank_in_group):
+ super(HaloExchangerSendRecv, self).__init__(ranks, rank_in_group)
+ nccl_id = inc.get_unique_nccl_id(1).cuda()
+ torch.distributed.broadcast(nccl_id, 0)
+ nccl_id = nccl_id.cpu()
+ print("%d :: nccl_id = %s" % (torch.distributed.get_rank(), str(nccl_id)))
+ # Create another global nccl communicator in addition to the one created by torch.distributed.init_process_group("nccl")
+ # This is unavoidable because the underlying NCCL communicator torch.distributed creates is a protected variable, hence
+ # it cannot be accessed from another class.
+ # TODO: Figure out a way to avoid creating a second global communicator
+ assert(torch.distributed.get_rank() == self.ranks[self.rank_in_group]), "ranks[%d](%d) != torch.distributed.get_rank()(%d)" % (self.rank_in_group, self.ranks[self.rank_in_group], torch.distributed.get_rank())
+ self.handle = inc.init_nccl_comm(nccl_id, torch.distributed.get_rank(), torch.distributed.get_world_size())
+
+ def left_right_halo_exchange(self, left_output_halo, right_output_halo, left_input_halo=None, right_input_halo=None):
+ if left_input_halo is None:
+ left_input_halo, right_input_halo = inc.left_right_halo_exchange(self.handle, self.left_rank, self.right_rank , left_output_halo, right_output_halo)
+ return left_input_halo, right_input_halo
+ else:
+ inc.left_right_halo_exchange_inplace(self.handle, self.left_rank, self.right_rank, left_output_halo, right_output_halo, left_input_halo, right_input_halo)
+
+class HaloExchangerPeer(HaloExchanger):
+ def __init__(self, ranks, rank_in_group, peer_pool, explicit_nhwc, numSM=0):
+ super(HaloExchangerPeer, self).__init__(ranks, rank_in_group)
+ self.diagnostics = False
+ self.explicit_nhwc = explicit_nhwc
+ self.numSM = numSM
+ self.peer_pool = peer_pool
+
+ def _allocate_peer_tensor(self, halo):
+
+ # Compute size in bytes
+ # Note: Pad buffer so each CUDA block gets required buffer size
+ size = 4 * halo.numel() * halo.element_size()
+ size_per_block = 128 * 2 * 16 # 128 threads each require two 128b buffers
+ size = (size + size_per_block - 1) // size_per_block * size_per_block
+
+ # Construct dtype peer buffer with desired size
+ shape = [1, 1, 1, size // halo.element_size()]
+ return self.peer_pool.allocate_peer_tensors(shape, halo.dtype, False, True)
+
+ def left_right_halo_exchange(self, left_output_halo, right_output_halo, left_input_halo=None, right_input_halo=None):
+ inplace = False if left_input_halo is None and right_input_halo is None else True
+ if not inplace:
+ left_input_halo = torch.empty_like(right_output_halo)
+ right_input_halo = torch.empty_like(left_output_halo)
+ channels_last = left_output_halo.is_contiguous(memory_format=torch.channels_last) and not self.explicit_nhwc
+ left_tx = self._allocate_peer_tensor(left_input_halo)
+ right_tx = self._allocate_peer_tensor(right_input_halo)
+ pm.push_pull_halos_1d(
+ self.diagnostics, self.explicit_nhwc, self.numSM, self.rank_in_group,
+ self.left_zero, left_output_halo, left_tx[self.rank_in_group], right_tx[self.wrap_around_left_rank_in_group], left_input_halo,
+ self.right_zero, right_output_halo, right_tx[self.rank_in_group], left_tx[self.wrap_around_right_rank_in_group], right_input_halo,
+ )
+ if not inplace:
+ return left_input_halo, right_input_halo
+
+# Class that combines input volume with halos from neighbors (1d).
+class HaloPadder:
+ def __init__(self, halo_ex):
+ self.halo_ex = halo_ex
+ self.stream1 = torch.cuda.Stream()
+ self.stream2 = torch.cuda.Stream()
+
+ def __call__(self, y, half_halo, explicit_nhwc, H_split):
+ channels_last = not explicit_nhwc and y.is_contiguous(memory_format=torch.channels_last)
+ if explicit_nhwc:
+ N,H,W,C = list(y.shape)
+ if H_split:
+ padded_shape = [N,H+2*half_halo,W,C]
+ ypad = torch.empty(shape=padded_shape, dtype=y.dtype, device=y.device, memory_format=torch.contiguous_format)
+ yleft = ypad[:,:half_halo,:,:]
+ ymid = ypad[:,half_halo:H+half_halo,:,:]
+ yright = ypad[:,H+half_halo:H+2*half_halo,:,:]
+ oleft = y[:,:half_halo,:,:]
+ oright = y[:,H-half_halo:,:,:]
+ else:
+ padded_shape = [N,H,W+2*half_halo,C]
+ ypad = torch.empty(shape=padded_shape, dtype=y.dtype, device=y.device, memory_format=torch.contiguous_format)
+ yleft = ypad[:,:,:half_halo,:]
+ ymid = ypad[:,:,half_halo:W+half_halo,:]
+ yright = ypad[:,:,W+half_halo:W+2*half_halo,:]
+ oleft = y[:,:,:half_halo,:]
+ oright = y[:,:,W-half_halo:,:]
+ else:
+ N,C,H,W = list(y.shape)
+ if H_split:
+ padded_shape = [N,C,H+2*half_halo,W]
+ ypad = torch.empty(shape=padded_shape, dtype=y.dtype, device=y.device, memory_format=torch.channels_last)
+ yleft = ypad[:,:,:half_halo,:]
+ ymid = ypad[:,:,half_halo:H+half_halo,:]
+ yright = ypad[:,:,H+half_halo:H+2*half_halo,:]
+ oleft = y[:,:,:half_halo,:]
+ oright = y[:,:,H-half_halo:,:]
+ else:
+ padded_shape = [N,C,H,W+2*half_halo]
+ ypad = torch.empty(shape=padded_shape, dtype=y.dtype, device=y.device, memory_format=torch.channels_last)
+ yleft = ypad[:,:,:,:half_halo]
+ ymid = ypad[:,:,:,half_halo:W+half_halo]
+ yright = ypad[:,:,:,W+half_halo:W+2*half_halo]
+ oleft = y[:,:,:,:half_halo]
+ oright = y[:,:,:,W-half_halo:]
+ with torch.cuda.stream(self.stream1):
+ self.halo_ex(oleft, oright, yleft, yright)
+ with torch.cuda.stream(self.stream2):
+ ymid.copy_(y)
+ return ypad
+
+ def wait(self):
+ current_stream = torch.cuda.current_stream()
+ current_stream.wait_stream(self.stream1)
+ current_stream.wait_stream(self.stream2)
diff --git a/apex/apex/contrib/bottleneck/test.py b/apex/apex/contrib/bottleneck/test.py
new file mode 100644
index 00000000..2c3c6213
--- /dev/null
+++ b/apex/apex/contrib/bottleneck/test.py
@@ -0,0 +1,71 @@
+import torch
+from bottleneck import Bottleneck
+torch.manual_seed(23337)
+
+# use True to print layerwise sum for all outputs in reference code path
+DEBUG = False#True
+
+for stride, o_channel in [(1,32), (1,128), (2,32)]:
+ print("testing stride ==", stride, ", in_channel == 32 , out_channel ==", o_channel)
+ a_ = torch.randn(17,32,28,28)
+
+ a = a_.cuda().half().to(memory_format=torch.channels_last).requires_grad_()
+ model = Bottleneck(32,8,o_channel,stride=stride).cuda().half().to(memory_format=torch.channels_last)
+
+ # test model
+ b = model(a)
+ b.mean().backward()
+ d_grad = a.grad.float()
+ a.grad = None
+ torch.cuda.synchronize()
+
+ if DEBUG:
+ print("[DEBUG] ref dx :", d_grad.sum().item())
+ # print wgrad. we don't need to reset since later cpp print before accumulation
+ for i, w in enumerate(model.w_conv):
+ print("[DEBUG] ref wgrad{} :".format(i+1), w.grad.sum().item())
+
+ wgrads = []
+ for w in model.w_conv:
+ wgrads.append(w.grad.float())
+
+ model.use_cudnn = True
+ model.zero_grad()
+ c = model(a)
+ c.mean().backward()
+
+ torch.cuda.synchronize()
+ print("comparing native and channels_last:")
+ print("max error fprop:", (b-c).abs().max().item(), "max elem:", b.abs().max().item())
+ print("max error dgrad:", (d_grad-a.grad.float()).abs().max().item(), "max elem:", d_grad.abs().max().item())
+ for i, (w, wgrad) in enumerate(zip(model.w_conv, wgrads)):
+ print("max error wgrad{}:".format(i+1), (wgrad - w.grad.float()).abs().max().item(), "max elem:", wgrad.abs().max().item())
+
+ nhwc_a = a_.permute(0,2,3,1).contiguous().cuda().half().requires_grad_()
+ nhwc_model = Bottleneck(32,8,o_channel,stride=stride,explicit_nhwc=True, use_cudnn=True).cuda().half()
+ for p,q in zip(model.parameters(), nhwc_model.parameters()):
+ # model's storage is already in nhwc, we clone and assign to explicit nhwc model
+ q.data.copy_(p.data.permute(0,2,3,1).contiguous())
+ for p,q in zip(model.buffers(), nhwc_model.buffers()):
+ q.data.copy_(p.data)
+
+ d = nhwc_model(nhwc_a)
+ d.mean().backward()
+ torch.cuda.synchronize()
+
+ # reset reference to cudnn channels_last permute
+ #c_s = c.storage().tolist()
+ #d_s = d.storage().tolist()
+ #print(max([x-y for x,y in zip(c_s,d_s)]))
+ c = c.contiguous(memory_format=torch.contiguous_format).permute(0,2,3,1).contiguous()
+ d_grad = a.grad.float().permute(0,2,3,1).contiguous()
+ wgrads = []
+ for w in model.w_conv:
+ wgrads.append(w.grad.float().permute(0,2,3,1).contiguous())
+
+ torch.cuda.synchronize()
+ print("comparing nhwc and channels_last:")
+ print("max error fprop:", (d-c).abs().max().item(), "max elem:", c.abs().max().item())
+ print("max error dgrad:", (d_grad-nhwc_a.grad.float()).abs().max().item(), "max elem:", d_grad.abs().max().item())
+ for i, (w, wgrad) in enumerate(zip(nhwc_model.w_conv, wgrads)):
+ print("max error wgrad{}:".format(i+1), (wgrad - w.grad.float()).abs().max().item(), "max elem:", wgrad.abs().max().item())
diff --git a/apex/apex/contrib/clip_grad/__init__.py b/apex/apex/contrib/clip_grad/__init__.py
new file mode 100644
index 00000000..cc9f501b
--- /dev/null
+++ b/apex/apex/contrib/clip_grad/__init__.py
@@ -0,0 +1 @@
+from .clip_grad import clip_grad_norm_
diff --git a/apex/apex/contrib/clip_grad/clip_grad.py b/apex/apex/contrib/clip_grad/clip_grad.py
new file mode 100644
index 00000000..b6411352
--- /dev/null
+++ b/apex/apex/contrib/clip_grad/clip_grad.py
@@ -0,0 +1,128 @@
+from typing import Union, Iterable
+
+import torch
+
+_kernel_import_succeeded = False
+try:
+ import amp_C
+ from apex.multi_tensor_apply import multi_tensor_applier
+ _kernel_import_succeeded = True
+except ImportError:
+ _kernel_import_succeeded = False
+
+_tensor_or_tensors = Union[torch.Tensor, Iterable[torch.Tensor]]
+
+
+def clip_grad_norm_(
+ parameters: _tensor_or_tensors, max_norm: float, norm_type: float = 2.0,
+ error_if_nonfinite: bool = False) -> torch.Tensor:
+ r"""Clips gradient norm of an iterable of parameters.
+
+ The norm is computed over all gradients together, as if they were
+ concatenated into a single vector. Gradients are modified in-place.
+
+ This is identical to torch.nn.utils.clip_grad_norm_, except it
+ uses a fused CUDA kernel when computing the 2-norm of GPU tensors
+ in float32 and float16.
+
+ Args:
+ parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a
+ single Tensor that will have gradients normalized
+ max_norm (float or int): max norm of the gradients
+ norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for
+ infinity norm.
+ error_if_nonfinite (bool): if True, an error is thrown if the total
+ norm of the gradients from :attr:`parameters` is ``nan``,
+ ``inf``, or ``-inf``. Default: False (will switch to True in the future)
+
+ Returns:
+ Total norm of the parameters (viewed as a single vector).
+
+ """
+ if isinstance(parameters, torch.Tensor):
+ parameters = [parameters]
+ parameters = [p for p in parameters if p.grad is not None]
+ max_norm = float(max_norm)
+ norm_type = float(norm_type)
+
+ # Trivial case
+ if len(parameters) == 0:
+ return torch.tensor(0.)
+
+ # Fallback implementation
+ if not (_kernel_import_succeeded
+ and norm_type == 2.0
+ and any(p.is_cuda for p in parameters)):
+ return torch.nn.utils.clip_grad_norm_(
+ parameters,
+ max_norm,
+ norm_type=norm_type,
+ error_if_nonfinite = error_if_nonfinite,
+ )
+
+ # Find fp32 and fp16 gradients on GPU
+ device = next(p.device for p in parameters if p.is_cuda)
+ grads_fp32, grads_fp16, grads_misc = [], [], []
+ for p in parameters:
+ grad = p.grad.detach()
+ if p.dtype == torch.float32 and p.device == device:
+ grads_fp32.append(grad)
+ elif p.dtype == torch.float16 and p.device == device:
+ grads_fp16.append(grad)
+ else:
+ grads_misc.append(grad)
+
+ # Compute gradient L2 norms
+ norms = []
+ dummy_overflow_buf = torch.zeros([1], dtype=torch.int32, device=device)
+ if grads_fp32:
+ norms.append(
+ multi_tensor_applier(
+ amp_C.multi_tensor_l2norm,
+ dummy_overflow_buf,
+ [grads_fp32],
+ False,
+ )[0]
+ )
+ if grads_fp16:
+ norms.append(
+ multi_tensor_applier(
+ amp_C.multi_tensor_l2norm,
+ dummy_overflow_buf,
+ [grads_fp16],
+ False,
+ )[0],
+ )
+ for g in grads_misc:
+ norms.append(torch.linalg.norm(g).unsqueeze(0).to(device))
+ total_norm = torch.linalg.norm(torch.cat(norms))
+
+ # Check for non-finite values
+ if error_if_nonfinite and torch.logical_or(total_norm.isnan(), total_norm.isinf()):
+ raise RuntimeError(
+ f'The total norm of order {norm_type} for gradients from '
+ '`parameters` is non-finite, so it cannot be clipped. To disable '
+ 'this error and scale the gradients by the non-finite norm anyway, '
+ 'set `error_if_nonfinite=False`')
+
+ # Scale gradients
+ clip_coef = max_norm / (total_norm + 1e-6)
+ clip_coef_clamped = torch.clamp(clip_coef, max=1.0)
+ if grads_fp32:
+ multi_tensor_applier(
+ amp_C.multi_tensor_scale,
+ dummy_overflow_buf,
+ [grads_fp32, grads_fp32],
+ clip_coef_clamped,
+ )
+ if grads_fp16:
+ multi_tensor_applier(
+ amp_C.multi_tensor_scale,
+ dummy_overflow_buf,
+ [grads_fp16, grads_fp16],
+ clip_coef_clamped,
+ )
+ for g in grads_misc:
+ g.mul_(clip_coef_clamped.to(g.device))
+
+ return total_norm
diff --git a/apex/apex/contrib/conv_bias_relu/__init__.py b/apex/apex/contrib/conv_bias_relu/__init__.py
new file mode 100644
index 00000000..6e449ba2
--- /dev/null
+++ b/apex/apex/contrib/conv_bias_relu/__init__.py
@@ -0,0 +1,2 @@
+from .conv_bias_relu import ConvBiasReLU, ConvBias, ConvBiasMaskReLU, ConvFrozenScaleBiasReLU
+
diff --git a/apex/apex/contrib/conv_bias_relu/conv_bias_relu.py b/apex/apex/contrib/conv_bias_relu/conv_bias_relu.py
new file mode 100644
index 00000000..c873ebe1
--- /dev/null
+++ b/apex/apex/contrib/conv_bias_relu/conv_bias_relu.py
@@ -0,0 +1,104 @@
+import pdb
+
+import torch
+from torch.autograd import gradcheck
+
+from apex import check_cudnn_version_and_warn
+import fused_conv_bias_relu
+
+check_cudnn_version_and_warn(__name__, 8400)
+
+
+class ConvBiasReLU_(torch.autograd.Function):
+ @staticmethod
+ @torch.cuda.amp.custom_fwd(cast_inputs=torch.half)
+ def forward(ctx, x, weight, bias, padding, stride):
+ outputs = fused_conv_bias_relu.forward([x, weight, bias], padding, stride)
+ ctx.save_for_backward(x, weight, outputs[0])
+ ctx.padding = padding
+ ctx.stride = stride
+
+ return outputs[0]
+
+ @staticmethod
+ @torch.cuda.amp.custom_bwd
+ def backward(ctx, grad_output):
+ bwd_args = [*ctx.saved_tensors, grad_output]
+ padding = ctx.padding
+ stride = ctx.stride
+ grads = fused_conv_bias_relu.backward(bwd_args, padding, stride)
+
+ return grads[0], grads[1], grads[2], None, None
+
+
+class ConvBiasMaskReLU_(torch.autograd.Function):
+ @staticmethod
+ @torch.cuda.amp.custom_fwd(cast_inputs=torch.half)
+ def forward(ctx, x, weight, bias, mask, padding, stride):
+ outputs = fused_conv_bias_relu.forward_mask([x, weight, bias, mask], padding, stride)
+ ctx.save_for_backward(x, weight, outputs[0])
+ ctx.padding = padding
+ ctx.stride = stride
+
+ return outputs[0]
+
+ @staticmethod
+ @torch.cuda.amp.custom_bwd
+ def backward(ctx, grad_output):
+ bwd_args = [*ctx.saved_tensors, grad_output]
+ padding = ctx.padding
+ stride = ctx.stride
+ grads = fused_conv_bias_relu.backward(bwd_args, padding, stride)
+
+ return grads[0], grads[1], grads[2], None, None, None
+
+
+class ConvBias_(torch.autograd.Function):
+ @staticmethod
+ @torch.cuda.amp.custom_fwd(cast_inputs=torch.half)
+ def forward(ctx, x, weight, bias, padding, stride):
+ outputs = fused_conv_bias_relu.forward_no_relu([x, weight, bias], padding, stride)
+ ctx.save_for_backward(x, weight)
+ ctx.padding = padding
+ ctx.stride = stride
+
+ return outputs[0]
+
+ @staticmethod
+ @torch.cuda.amp.custom_bwd
+ def backward(ctx, grad_output):
+ bwd_args = [*ctx.saved_tensors, grad_output]
+ padding = ctx.padding
+ stride = ctx.stride
+ grads = fused_conv_bias_relu.backward_no_relu(bwd_args, padding, stride)
+
+ return grads[0], grads[1], grads[2], None, None
+
+
+class ConvFrozenScaleBiasReLU_(torch.autograd.Function):
+ @staticmethod
+ @torch.cuda.amp.custom_fwd(cast_inputs=torch.half)
+ def forward(ctx, x, weight, scale, bias, padding, stride):
+ output = fused_conv_bias_relu.forward_cscale_cbias_relu([x, weight, scale, bias], padding, stride)
+ ctx.save_for_backward(x, weight, scale, output)
+ ctx.padding = padding
+ ctx.stride = stride
+
+ return output
+
+ @staticmethod
+ @torch.cuda.amp.custom_bwd
+ def backward(ctx, grad_output):
+ bwd_args = [*ctx.saved_tensors, grad_output]
+ padding = ctx.padding
+ stride = ctx.stride
+ grads = fused_conv_bias_relu.backward_cscale_cbias_relu(bwd_args, padding, stride)
+
+ return grads[0], grads[1], None, None, None, None
+
+
+ConvBiasReLU = ConvBiasReLU_.apply
+ConvBiasMaskReLU = ConvBiasMaskReLU_.apply
+ConvBias = ConvBias_.apply
+ConvFrozenScaleBiasReLU = ConvFrozenScaleBiasReLU_.apply
+
diff --git a/apex/apex/contrib/csrc/bottleneck/bottleneck.cpp b/apex/apex/contrib/csrc/bottleneck/bottleneck.cpp
new file mode 100644
index 00000000..9a0c3403
--- /dev/null
+++ b/apex/apex/contrib/csrc/bottleneck/bottleneck.cpp
@@ -0,0 +1,4073 @@
+#include
+#include // for getcudnnhandle
+#include
+#include
+#include
+#include
+
+#include
+
+#ifdef DEBUG
+#define DEBUG_MSG(str) do { std::cout << str << std::endl; } while( false )
+#else
+#define DEBUG_MSG(str) do { } while ( false )
+#endif
+
+#ifdef DEBUG_CUDNN
+#define DEBUG_CUDNN_MSG(buf, str) do { buf << str << std::endl; } while( false )
+#else
+#define DEBUG_CUDNN_MSG(buf, str) do { } while ( false )
+#endif
+
+#define checkCudnnErr(...) \
+ do { \
+ int err = checkCudnnError(__VA_ARGS__, #__VA_ARGS__, __FILE__, __LINE__); \
+ if (err) { \
+ return; \
+ } \
+ } while (0)
+
+
+int checkCudnnError(cudnnStatus_t code, const char* expr, const char* file, int line) {
+ if (code) {
+ printf("CUDNN error at %s:%d, code=%d (%s) in '%s'\n", file, line, (int)code, cudnnGetErrorString(code), expr);
+ return 1;
+ }
+ return 0;
+}
+
+void checkError(cudaError_t code, char const * func, const char *file, const int line, bool abort = true);
+#define checkCUDAError(val) { checkError((val), #val, __FILE__, __LINE__); } // in-line regular function
+
+void checkError(cudaError_t code, char const * func, const char *file, const int line, bool abort)
+{
+ if (code != cudaSuccess)
+ {
+ const char * errorMessage = cudaGetErrorString(code);
+ fprintf(stderr, "CUDA error returned from \"%s\" at %s:%d, Error code: %d (%s)\n", func, file, line, code, errorMessage);
+ if (abort){
+ cudaDeviceReset();
+ exit(code);
+ }
+ }
+}
+
+void generateStrides(const int64_t* dimA, int64_t* strideA, int nbDims, cudnnTensorFormat_t filterFormat) {
+ // For INT8x4 and INT8x32 we still compute standard strides here to input
+ // into the cuDNN functions. We will manually scale by resizeFactor in the cpu ref.
+ if (filterFormat == CUDNN_TENSOR_NCHW) {
+ strideA[nbDims - 1] = 1;
+ for (int64_t d = nbDims - 2; d >= 0; d--) {
+ strideA[d] = strideA[d + 1] * dimA[d + 1];
+ }
+ } else {
+ // Here we assume that the format is CUDNN_TENSOR_NHWC
+ strideA[1] = 1;
+ strideA[nbDims - 1] = strideA[1] * dimA[1];
+ for (int64_t d = nbDims - 2; d >= 2; d--) {
+ strideA[d] = strideA[d + 1] * dimA[d + 1];
+ }
+ strideA[0] = strideA[2] * dimA[2];
+ }
+}
+
+
+int getFwdConvDilatedFilterDim(int filterDim, int dilation) {
+ return ((filterDim - 1) * dilation) + 1;
+}
+
+int getFwdConvPaddedImageDim(int tensorDim, int pad) {
+ return tensorDim + (2 * pad);
+}
+
+int getFwdConvOutputDim(
+ int tensorDim,
+ int pad,
+ int filterDim,
+ int stride,
+ int dilation)
+{
+ int p = (getFwdConvPaddedImageDim(tensorDim, pad) - getFwdConvDilatedFilterDim(filterDim, dilation)) / stride + 1;
+ return (p);
+}
+
+enum {
+ X_TENSOR,
+ Y_TENSOR,
+ W_TENSOR,
+ Z_TENSOR,
+ B_TENSOR,
+ AFTERADD_TENSOR,
+ AFTERBIAS_TENSOR,
+ AFTERCONV_TENSOR,
+ OPTIONAL,
+ AFTEROPT_TENSOR,
+};
+
+using common_conv_descriptors =
+ std::tuple;
+
+
+common_conv_descriptors
+create_common_descriptors(int64_t* x_dim_padded,
+ int64_t* padA,
+ int64_t* convstrideA,
+ int64_t* dilationA,
+ int64_t* w_dim_padded,
+ int64_t* y_dim_padded,
+ cudnnDataType_t dataType,
+ cudnnConvolutionMode_t mode) {
+ const int convDim = 2;
+
+ int64_t strideA_padded[4];
+ int64_t outstrideA_padded[4];
+ int64_t filterstrideA_padded[4];
+
+ generateStrides(w_dim_padded, filterstrideA_padded, 4, CUDNN_TENSOR_NHWC);
+ generateStrides(x_dim_padded, strideA_padded, 4, CUDNN_TENSOR_NHWC);
+ generateStrides(y_dim_padded, outstrideA_padded, 4, CUDNN_TENSOR_NHWC);
+
+ return common_conv_descriptors(cudnn_frontend::TensorBuilder()
+ .setDim(4, x_dim_padded)
+ .setStrides(4, strideA_padded)
+ .setId('x')
+ .setAlignment(16)
+ .setDataType(dataType)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, y_dim_padded)
+ .setStrides(4, outstrideA_padded)
+ .setId('y')
+ .setAlignment(16)
+ .setDataType(dataType)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, w_dim_padded)
+ .setStrides(4, filterstrideA_padded)
+ .setId('w')
+ .setAlignment(16)
+ .setDataType(dataType)
+ .build(),
+ cudnn_frontend::ConvDescBuilder()
+ .setDataType(CUDNN_DATA_FLOAT)
+ .setMathMode(mode)
+ .setNDims(convDim)
+ .setStrides(convDim, convstrideA)
+ .setPrePadding(convDim, padA)
+ .setPostPadding(convDim, padA)
+ .setDilation(convDim, dilationA)
+ .build());
+}
+
+using common_convbias_descriptors = std::tuple;
+
+common_convbias_descriptors
+create_conv_bias_add_act_descriptors(int64_t* x_dim_padded,
+ int64_t* padA,
+ int64_t* convstrideA,
+ int64_t* dilationA,
+ int64_t* w_dim_padded,
+ int64_t* y_dim_padded,
+ cudnnDataType_t dataType) {
+ const int convDim = 2;
+
+ int64_t b_dim_padded[4];
+ b_dim_padded[0] = 1;
+ b_dim_padded[1] = y_dim_padded[1];
+ b_dim_padded[2] = 1;
+ b_dim_padded[3] = 1;
+
+ int64_t x_stride_padded[4];
+ int64_t y_stride_padded[4];
+ int64_t w_stride_padded[4];
+ int64_t b_stride_padded[4];
+
+ generateStrides(w_dim_padded, w_stride_padded, 4, CUDNN_TENSOR_NHWC);
+ generateStrides(x_dim_padded, x_stride_padded, 4, CUDNN_TENSOR_NHWC);
+ generateStrides(y_dim_padded, y_stride_padded, 4, CUDNN_TENSOR_NHWC);
+ generateStrides(b_dim_padded, b_stride_padded, 4, CUDNN_TENSOR_NHWC);
+
+ return common_convbias_descriptors(cudnn_frontend::TensorBuilder()
+ .setDim(4, x_dim_padded)
+ .setStrides(4, x_stride_padded)
+ .setId('x')
+ .setAlignment(16)
+ .setDataType(dataType)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, y_dim_padded)
+ .setStrides(4, y_stride_padded)
+ .setId('y')
+ .setAlignment(16)
+ .setDataType(dataType)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, w_dim_padded)
+ .setStrides(4, w_stride_padded)
+ .setId('w')
+ .setAlignment(16)
+ .setDataType(dataType)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, b_dim_padded)
+ .setStrides(4, b_stride_padded)
+ .setId('z')
+ .setAlignment(16)
+ .setDataType(dataType)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, b_dim_padded)
+ .setStrides(4, b_stride_padded)
+ .setId('b')
+ .setAlignment(16)
+ .setDataType(dataType)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, y_dim_padded)
+ .setStrides(4, y_stride_padded)
+ .setVirtual()
+ .setId('A') // after add
+ .setAlignment(16)
+ .setDataType(CUDNN_DATA_FLOAT)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, y_dim_padded)
+ .setStrides(4, y_stride_padded)
+ .setVirtual()
+ .setId('B') // after bias
+ .setAlignment(16)
+ .setDataType(CUDNN_DATA_FLOAT)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, y_dim_padded)
+ .setStrides(4, y_stride_padded)
+ .setId('C') // after conv
+ .setAlignment(16)
+ .setVirtual()
+ .setDataType(CUDNN_DATA_FLOAT)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, y_dim_padded)
+ .setStrides(4, y_stride_padded)
+ .setId('i')
+ .setAlignment(16)
+ .setDataType(dataType)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, y_dim_padded)
+ .setStrides(4, y_stride_padded)
+ .setId('D') // after optional add
+ .setAlignment(16)
+ .setVirtual()
+ .setDataType(CUDNN_DATA_FLOAT)
+ .build());
+}
+
+// tensor descriptors used for dgrad
+enum {
+ X_OR_DX_TENSOR,
+ DY_TENSOR,
+ W_OR_DW_TENSOR,
+ SCALE_TENSOR,
+ RELU_TENSOR,
+ AFTER_DCONV_TENSOR,
+ AFTER_DRELU_TENSOR,
+};
+
+using dconv_descriptors = std::tuple;
+
+dconv_descriptors
+create_dconv_descriptors(int64_t* x_dim_padded,
+ int64_t* padA,
+ int64_t* convstrideA,
+ int64_t* dilationA,
+ int64_t* w_dim_padded,
+ int64_t* y_dim_padded,
+ cudnnDataType_t dataType) {
+ const int convDim = 2;
+
+ int64_t b_dim_padded[4];
+ b_dim_padded[0] = 1;
+ b_dim_padded[1] = x_dim_padded[1];
+ b_dim_padded[2] = 1;
+ b_dim_padded[3] = 1;
+
+ int64_t x_stride_padded[4];
+ int64_t y_stride_padded[4];
+ int64_t w_stride_padded[4];
+ int64_t b_stride_padded[4];
+
+ generateStrides(w_dim_padded, w_stride_padded, 4, CUDNN_TENSOR_NHWC);
+ generateStrides(x_dim_padded, x_stride_padded, 4, CUDNN_TENSOR_NHWC);
+ generateStrides(y_dim_padded, y_stride_padded, 4, CUDNN_TENSOR_NHWC);
+ generateStrides(b_dim_padded, b_stride_padded, 4, CUDNN_TENSOR_NHWC);
+
+ return dconv_descriptors(cudnn_frontend::TensorBuilder()
+ .setDim(4, x_dim_padded)
+ .setStrides(4, x_stride_padded)
+ .setId('x')
+ .setAlignment(16)
+ .setDataType(dataType)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, y_dim_padded)
+ .setStrides(4, y_stride_padded)
+ .setId('y')
+ .setAlignment(16)
+ .setDataType(dataType)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, w_dim_padded)
+ .setStrides(4, w_stride_padded)
+ .setId('w')
+ .setAlignment(16)
+ .setDataType(dataType)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, b_dim_padded)
+ .setStrides(4, b_stride_padded)
+ .setId('s')
+ .setAlignment(16)
+ .setDataType(dataType)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, x_dim_padded)
+ .setStrides(4, x_stride_padded)
+ .setId('r')
+ .setAlignment(16)
+ .setDataType(dataType)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, x_dim_padded)
+ .setStrides(4, x_stride_padded)
+ .setVirtual()
+ .setId('A') // after dconv
+ .setAlignment(16)
+ .setDataType(CUDNN_DATA_FLOAT)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, x_dim_padded)
+ .setStrides(4, x_stride_padded)
+ .setVirtual()
+ .setId('B') // after drelu
+ .setAlignment(16)
+ .setDataType(CUDNN_DATA_FLOAT)
+ .build());
+}
+
+// create a cache for plan
+std::unordered_map plan_cache;
+
+// TODO: better name
+std::string getConvFusionString(int64_t* x_dim_padded,
+ int64_t* padA,
+ int64_t* convstrideA,
+ int64_t* dilationA,
+ int64_t* w_dim_padded,
+ cudnnDataType_t dataType,
+ std::string fusion_string) {
+
+ for(int i=0;i<4;i++) {
+ fusion_string += 'X';
+ fusion_string += std::to_string(x_dim_padded[i]);
+ }
+ for(int i=0;i<4;i++) {
+ fusion_string += 'W';
+ fusion_string += std::to_string(w_dim_padded[i]);
+ }
+ for(int i=0;i<2;i++) {
+ fusion_string += 'P';
+ fusion_string += std::to_string(padA[i]);
+ }
+ for(int i=0;i<2;i++) {
+ fusion_string += 'S';
+ fusion_string += std::to_string(convstrideA[i]);
+ }
+ for(int i=0;i<2;i++) {
+ fusion_string += 'D';
+ fusion_string += std::to_string(dilationA[i]);
+ }
+ fusion_string += 'T';
+ fusion_string += std::to_string(dataType);
+ return fusion_string;
+}
+
+cudnn_frontend::ExecutionPlan& getOrCreatePlan(cudnnHandle_t handle_,
+ std::stringstream& log_buf,
+ cudnn_frontend::OperationGraph& opGraph,
+ std::string cache_string,
+ bool use_heuristic = true){
+ auto it = plan_cache.find(cache_string);
+ if (it != plan_cache.end()) {
+ DEBUG_CUDNN_MSG(log_buf, "Found plan in cache");
+ return it->second;
+ } else {
+ if (use_heuristic){
+ // TODO: confirm which mode to use
+ auto heuristics = cudnn_frontend::EngineHeuristicsBuilder()
+ .setOperationGraph(opGraph)
+ .setHeurMode(CUDNN_HEUR_MODE_INSTANT)
+ .build();
+ // try 3 times for now as WAR for no heuristic training
+ int max_tries = 3, count = 0;
+ auto& engine_configs = heuristics.getEngineConfig(max_tries);
+ while(true) {
+ try {
+ plan_cache.emplace(cache_string, std::move(cudnn_frontend::ExecutionPlanBuilder()
+ .setHandle(handle_)
+ .setEngineConfig(engine_configs[count], opGraph.getTag())
+ .build()));
+ break;
+ } catch (cudnn_frontend::cudnnException e) {
+ if (++count == max_tries) throw e;
+ }
+ }
+ }else{
+ DEBUG_CUDNN_MSG(log_buf, "No plan in cache");
+ // How many engines support this operation graph ?
+ auto total_engines = opGraph.getEngineCount();
+ DEBUG_CUDNN_MSG(log_buf, opGraph.describe() << " has " << total_engines << " engines.");
+ // We have to randomly pick one engine from [0, total_engines)
+ // Selecting "0" by default
+ auto engine = cudnn_frontend::EngineBuilder().setGlobalEngineIdx(0).setOperationGraph(opGraph).build();
+ DEBUG_CUDNN_MSG(log_buf, engine.describe());
+ auto& knobs = engine.getSupportedKnobs();
+ for (auto it = std::begin(knobs); it != std::end(knobs); ++it) {
+ DEBUG_CUDNN_MSG(log_buf, it->describe());
+ }
+ if (knobs.begin() != knobs.end()) {
+ DEBUG_CUDNN_MSG(log_buf, "Updated knob choice");
+ knobs.begin()->setChoice(knobs.begin()->getMinValue() + 1);
+ DEBUG_CUDNN_MSG(log_buf, knobs.begin()->describe());
+ }
+
+ // Createmplacee the requisite engine config
+ auto engine_config = cudnn_frontend::EngineConfigBuilder().setEngine(engine).build();
+ DEBUG_CUDNN_MSG(log_buf, engine_config.describe());
+ plan_cache.emplace(cache_string, std::move(cudnn_frontend::ExecutionPlanBuilder().setHandle(handle_).setEngineConfig(engine_config).build()));
+ }
+
+ return plan_cache.find(cache_string)->second;
+ }
+}
+
+void
+run_conv_scale_bias_add_activation(int64_t* x_dim_padded,
+ int64_t* pad,
+ int64_t* convstride,
+ int64_t* dilation,
+ int64_t* w_dim_padded,
+ int64_t* y_dim_padded,
+ cudnnDataType_t dataType,
+ at::Half* devPtrX,
+ at::Half* devPtrW,
+ at::Half* devPtrY,
+ at::Half* devPtrZ,
+ at::Half* devPtrB,
+ at::Half* devPtrI) {
+ cudnnHandle_t handle_ = torch::native::getCudnnHandle();
+ std::stringstream log_buf;
+ try {
+ int convDim = 2;
+
+ // Creates the necessary tensor descriptors
+ common_convbias_descriptors tensors = create_conv_bias_add_act_descriptors(
+ x_dim_padded, pad, convstride, dilation, w_dim_padded, y_dim_padded, dataType);
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+
+ // Define the add operation
+ auto scaleDesc = cudnn_frontend::PointWiseDescBuilder()
+ .setMode(CUDNN_POINTWISE_MUL)
+ .setMathPrecision(CUDNN_DATA_FLOAT)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, scaleDesc.describe());
+
+ // Define the bias operation
+ auto biasDesc = cudnn_frontend::PointWiseDescBuilder()
+ .setMode(CUDNN_POINTWISE_ADD)
+ .setMathPrecision(CUDNN_DATA_FLOAT)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, biasDesc.describe());
+
+ // optional add
+ auto addDesc = cudnn_frontend::PointWiseDescBuilder()
+ .setMode(CUDNN_POINTWISE_ADD)
+ .setMathPrecision(CUDNN_DATA_FLOAT)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, addDesc.describe());
+
+ // Define the activation operation
+ auto actDesc = cudnn_frontend::PointWiseDescBuilder()
+ .setMode(CUDNN_POINTWISE_RELU_FWD)
+ .setMathPrecision(CUDNN_DATA_FLOAT)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, actDesc.describe());
+
+ // Define the convolution problem
+ auto convDesc = cudnn_frontend::ConvDescBuilder()
+ .setDataType(CUDNN_DATA_FLOAT)
+ .setMathMode(CUDNN_CROSS_CORRELATION)
+ .setNDims(convDim)
+ .setStrides(convDim, convstride)
+ .setPrePadding(convDim, pad)
+ .setPostPadding(convDim, pad)
+ .setDilation(convDim, dilation)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, convDesc.describe());
+
+ float alpha = 1.0f;
+ float beta = 0.0f;
+
+ // Create a convolution Node
+ auto conv_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_CONVOLUTION_FORWARD_DESCRIPTOR)
+ .setxDesc(std::get(tensors))
+ .setwDesc(std::get(tensors))
+ .setyDesc(std::get(tensors))
+ .setcDesc(convDesc)
+ .setAlpha(alpha)
+ .setBeta(beta)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, conv_op.describe());
+
+ // Create a Add Node with scaling parameters.
+ auto scale_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)
+ .setxDesc(conv_op.getOutputTensor())
+ .setbDesc(std::get(tensors))
+ .setyDesc(std::get(tensors))
+ .setpwDesc(scaleDesc)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, scale_op.describe());
+
+ // Create a Bias Node.
+ auto bias_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)
+ .setxDesc(scale_op.getOutputTensor())
+ .setbDesc(std::get(tensors))
+ .setyDesc(std::get(tensors))
+ .setpwDesc(biasDesc)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, bias_op.describe());
+
+ // Create a optional add Node.
+ auto add_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)
+ .setxDesc(bias_op.getOutputTensor())
+ .setbDesc(std::get(tensors))
+ .setyDesc(std::get(tensors))
+ .setpwDesc(addDesc)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, add_op.describe());
+
+
+ // Create an Activation Node.
+ auto act_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)
+ .setxDesc(devPtrI ? add_op.getOutputTensor() : bias_op.getOutputTensor())
+ .setyDesc(std::get(tensors))
+ .setpwDesc(actDesc)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, act_op.describe());
+
+ // Create an Operation Graph. In this case it is convolution add bias activation
+ std::array ops = {&conv_op, &scale_op, &bias_op, devPtrI ? &add_op : &act_op, &act_op};
+
+ auto opGraph = cudnn_frontend::OperationGraphBuilder()
+ .setHandle(handle_)
+ .setOperationGraph(devPtrI ? ops.size() : 4, ops.data())
+ .build();
+
+ // Create string encoding for plan caching
+ auto cache_string = getConvFusionString(x_dim_padded, pad, convstride, dilation, w_dim_padded, dataType, opGraph.getTag());
+ DEBUG_CUDNN_MSG(log_buf, "[convstring] " << cache_string);
+
+ auto& plan = getOrCreatePlan(handle_, log_buf, opGraph, cache_string);
+ DEBUG_CUDNN_MSG(log_buf, "Plan tag: " << plan.getTag());
+
+ auto workspace_size = plan.getWorkspaceSize();
+ DEBUG_CUDNN_MSG(log_buf, plan.describe() << " requires workspace " << workspace_size);
+
+ void* workspace_ptr = nullptr;
+ auto workspace_tensor = at::empty({(workspace_size+3)/4}, at::TensorOptions(at::kCUDA).dtype(at::kFloat));
+ if (workspace_size > 0) {
+ workspace_ptr = workspace_tensor.data_ptr();
+ }
+ void* data_ptrs[] = {devPtrX, devPtrY, devPtrW, devPtrZ, devPtrB, devPtrI};
+ int64_t uids[] = {'x', 'y', 'w', 'z', 'b', 'i'};
+ auto variantPack = cudnn_frontend::VariantPackBuilder()
+ .setWorkspacePointer(workspace_ptr)
+ .setDataPointers(devPtrI ? 6 : 5, data_ptrs)
+ .setUids(devPtrI ? 6 : 5, uids)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, "variantPack " << variantPack.describe());
+ cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc());
+ checkCudnnErr(status);
+ cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, "Plan execute error", status);
+ } catch (cudnn_frontend::cudnnException e) {
+ std::cout << log_buf.str() << "[ERROR] Exception " << e.what() << std::endl;
+ }
+}
+
+void
+run_conv_scale_bias(int64_t* x_dim_padded,
+ int64_t* pad,
+ int64_t* convstride,
+ int64_t* dilation,
+ int64_t* w_dim_padded,
+ int64_t* y_dim_padded,
+ cudnnDataType_t dataType,
+ at::Half* devPtrX,
+ at::Half* devPtrW,
+ at::Half* devPtrY,
+ at::Half* devPtrZ,
+ at::Half* devPtrB) {
+ cudnnHandle_t handle_ = torch::native::getCudnnHandle();
+ std::stringstream log_buf;
+ try {
+ int convDim = 2;
+
+ // Creates the necessary tensor descriptors
+ common_convbias_descriptors tensors = create_conv_bias_add_act_descriptors(
+ x_dim_padded, pad, convstride, dilation, w_dim_padded, y_dim_padded, dataType);
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+
+ // Define the add operation
+ auto scaleDesc = cudnn_frontend::PointWiseDescBuilder()
+ .setMode(CUDNN_POINTWISE_MUL)
+ .setMathPrecision(CUDNN_DATA_FLOAT)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, scaleDesc.describe());
+
+ // Define the bias operation
+ auto addDesc = cudnn_frontend::PointWiseDescBuilder()
+ .setMode(CUDNN_POINTWISE_ADD)
+ .setMathPrecision(CUDNN_DATA_FLOAT)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, addDesc.describe());
+
+ // Define the convolution problem
+ auto convDesc = cudnn_frontend::ConvDescBuilder()
+ .setDataType(CUDNN_DATA_FLOAT)
+ .setMathMode(CUDNN_CROSS_CORRELATION)
+ .setNDims(convDim)
+ .setStrides(convDim, convstride)
+ .setPrePadding(convDim, pad)
+ .setPostPadding(convDim, pad)
+ .setDilation(convDim, dilation)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, convDesc.describe());
+
+ float alpha = 1.0f;
+ float beta = 0.0f;
+
+ // Create a convolution Node
+ auto conv_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_CONVOLUTION_FORWARD_DESCRIPTOR)
+ .setxDesc(std::get(tensors))
+ .setwDesc(std::get(tensors))
+ .setyDesc(std::get(tensors))
+ .setcDesc(convDesc)
+ .setAlpha(alpha)
+ .setBeta(beta)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, conv_op.describe());
+
+ // Create a Add Node with scaling parameters.
+ auto scale_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)
+ .setxDesc(conv_op.getOutputTensor())
+ .setbDesc(std::get(tensors))
+ .setyDesc(std::get(tensors)) // TODO: change enum to aftermul
+ .setpwDesc(scaleDesc)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, scale_op.describe());
+
+ // Create a Bias Node.
+ auto add_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)
+ .setxDesc(scale_op.getOutputTensor())
+ .setbDesc(std::get(tensors))
+ .setyDesc(std::get(tensors))
+ .setpwDesc(addDesc)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, add_op.describe());
+
+ // Create an Operation Graph. In this case it is convolution add bias activation
+ std::array ops = {&conv_op, &scale_op, &add_op};
+
+ auto opGraph = cudnn_frontend::OperationGraphBuilder()
+ .setHandle(handle_)
+ .setOperationGraph(ops.size(), ops.data())
+ .build();
+
+ // Create string encoding for plan caching
+ auto cache_string = getConvFusionString(x_dim_padded, pad, convstride, dilation, w_dim_padded, dataType, opGraph.getTag());
+ DEBUG_CUDNN_MSG(log_buf, "[convstring] " << cache_string);
+
+ auto& plan = getOrCreatePlan(handle_, log_buf, opGraph, cache_string);
+ DEBUG_CUDNN_MSG(log_buf, "Plan tag: " << plan.getTag());
+
+ auto workspace_size = plan.getWorkspaceSize();
+ DEBUG_CUDNN_MSG(log_buf, plan.describe() << " requires workspace " << workspace_size);
+
+ void* workspace_ptr = nullptr;
+ auto workspace_tensor = at::empty({(workspace_size+3)/4}, at::TensorOptions(at::kCUDA).dtype(at::kFloat));
+ if (workspace_size > 0) {
+ workspace_ptr = workspace_tensor.data_ptr();
+ }
+ void* data_ptrs[] = {devPtrX, devPtrY, devPtrW, devPtrZ, devPtrB};
+ int64_t uids[] = {'x', 'y', 'w', 'z', 'b'};
+ auto variantPack = cudnn_frontend::VariantPackBuilder()
+ .setWorkspacePointer(workspace_ptr)
+ .setDataPointers(5, data_ptrs)
+ .setUids(5, uids)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, "variantPack " << variantPack.describe());
+ cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc());
+ checkCudnnErr(status);
+ cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, "Plan execute error", status);
+ } catch (cudnn_frontend::cudnnException e) {
+ std::cout << log_buf.str() << "[ERROR] Exception " << e.what() << std::endl;
+ }
+}
+
+
+void
+run_dconv_drelu_dscale(int64_t* x_dim_padded,
+ int64_t* pad,
+ int64_t* convstride,
+ int64_t* dilation,
+ int64_t* w_dim_padded,
+ int64_t* y_dim_padded,
+ cudnnDataType_t dataType,
+ at::Half* devPtrX,
+ at::Half* devPtrW,
+ at::Half* devPtrY,
+ at::Half* devPtrZ,
+ at::Half* devPtrR) {
+ cudnnHandle_t handle_ = torch::native::getCudnnHandle();
+ std::stringstream log_buf;
+ try {
+ int convDim = 2;
+
+ // Creates the necessary tensor descriptors
+ dconv_descriptors tensors = create_dconv_descriptors(
+ x_dim_padded, pad, convstride, dilation, w_dim_padded, y_dim_padded, dataType);
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+
+ // Define the convolution problem
+ auto convDesc = cudnn_frontend::ConvDescBuilder()
+ .setDataType(CUDNN_DATA_FLOAT)
+ .setMathMode(CUDNN_CROSS_CORRELATION)
+ .setNDims(convDim)
+ .setStrides(convDim, convstride)
+ .setPrePadding(convDim, pad)
+ .setPostPadding(convDim, pad)
+ .setDilation(convDim, dilation)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, convDesc.describe());
+
+ // Define the activation backward operation
+ auto actDesc = cudnn_frontend::PointWiseDescBuilder()
+ .setMode(CUDNN_POINTWISE_RELU_BWD)
+ .setMathPrecision(CUDNN_DATA_FLOAT)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, actDesc.describe());
+
+ // Define the scale backward operation
+ auto scaleDesc = cudnn_frontend::PointWiseDescBuilder()
+ .setMode(CUDNN_POINTWISE_MUL)
+ .setMathPrecision(CUDNN_DATA_FLOAT)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, scaleDesc.describe());
+
+ float alpha = 1.0f;
+ float beta = 0.0f;
+
+ // Create a convolution Node
+ auto conv_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_DATA_DESCRIPTOR)
+ .setdxDesc(std::get(tensors))
+ .setwDesc(std::get(tensors))
+ .setdyDesc(std::get(tensors))
+ .setcDesc(convDesc)
+ .setAlpha(alpha)
+ .setBeta(beta)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, conv_op.describe());
+
+ // TODO: do we need getOutputTensor(), and what it returns in backward case?
+ // Create an relu backward Node.
+ auto act_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)
+ .setdyDesc(std::get(tensors))
+ .setxDesc(std::get(tensors))
+ .setdxDesc(std::get(tensors))
+ .setpwDesc(actDesc)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, act_op.describe());
+
+ // Create a Scale Node.
+ auto scale_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)
+ .setxDesc(std::get(tensors))
+ .setbDesc(std::get(tensors))
+ .setyDesc(std::get(tensors))
+ .setpwDesc(scaleDesc)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, scale_op.describe());
+
+ // Create an Operation Graph. In this case it is convolution add bias activation
+ std::array ops = {&conv_op, &act_op, &scale_op};
+
+ auto opGraph = cudnn_frontend::OperationGraphBuilder()
+ .setHandle(handle_)
+ .setOperationGraph(ops.size(), ops.data())
+ .build();
+
+ // Create string encoding for plan caching
+ auto cache_string = getConvFusionString(x_dim_padded, pad, convstride, dilation, w_dim_padded, dataType, opGraph.getTag());
+ DEBUG_CUDNN_MSG(log_buf, "[convstring] " << cache_string);
+
+ auto& plan = getOrCreatePlan(handle_, log_buf, opGraph, cache_string);
+ DEBUG_CUDNN_MSG(log_buf, "Plan tag: " << plan.getTag());
+
+ auto workspace_size = plan.getWorkspaceSize();
+ DEBUG_CUDNN_MSG(log_buf, plan.describe() << " requires workspace " << workspace_size);
+
+ void* workspace_ptr = nullptr;
+ auto workspace_tensor = at::empty({(workspace_size+3)/4}, at::TensorOptions(at::kCUDA).dtype(at::kFloat));
+ if (workspace_size > 0) {
+ workspace_ptr = workspace_tensor.data_ptr();
+ }
+ void* data_ptrs[] = {devPtrX, devPtrY, devPtrW, devPtrZ, devPtrR};
+ int64_t uids[] = {'x', 'y', 'w', 's', 'r'};
+ auto variantPack = cudnn_frontend::VariantPackBuilder()
+ .setWorkspacePointer(workspace_ptr)
+ .setDataPointers(5, data_ptrs)
+ .setUids(5, uids)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, "variantPack " << variantPack.describe());
+ cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc());
+ checkCudnnErr(status);
+ cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, "Plan execute error", status);
+ } catch (cudnn_frontend::cudnnException e) {
+ std::cout << log_buf.str() << "[ERROR] Exception " << e.what() << std::endl;
+ }
+}
+
+void
+run_dconv(int64_t* x_dim_padded,
+ int64_t* pad,
+ int64_t* convstride,
+ int64_t* dilation,
+ int64_t* w_dim_padded,
+ int64_t* y_dim_padded,
+ cudnnDataType_t dataType,
+ at::Half* devPtrX,
+ at::Half* devPtrW,
+ at::Half* devPtrY,
+ cudnnBackendDescriptorType_t mode) {
+ cudnnHandle_t handle_ = torch::native::getCudnnHandle();
+ std::stringstream log_buf;
+ try {
+ int convDim = 2;
+
+ // Creates the necessary tensor descriptors
+ dconv_descriptors tensors = create_dconv_descriptors(
+ x_dim_padded, pad, convstride, dilation, w_dim_padded, y_dim_padded, dataType);
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+
+ // Define the convolution problem
+ auto convDesc = cudnn_frontend::ConvDescBuilder()
+ .setDataType(CUDNN_DATA_FLOAT)
+ .setMathMode(CUDNN_CROSS_CORRELATION)
+ .setNDims(convDim)
+ .setStrides(convDim, convstride)
+ .setPrePadding(convDim, pad)
+ .setPostPadding(convDim, pad)
+ .setDilation(convDim, dilation)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, convDesc.describe());
+
+ float alpha = 1.0f;
+ float beta = 0.0f;
+
+ // Create a convolution Node
+ // mode should be one of following
+ // CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_DATA_DESCRIPTOR
+ // CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_FILTER_DESCRIPTOR
+ auto conv_op_builder = cudnn_frontend::OperationBuilder(mode);
+ if (mode == CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_DATA_DESCRIPTOR) {
+ conv_op_builder.setdxDesc(std::get(tensors))
+ .setwDesc(std::get(tensors))
+ .setdyDesc(std::get(tensors))
+ .setcDesc(convDesc)
+ .setAlpha(alpha)
+ .setBeta(beta);
+ }
+ else {
+ conv_op_builder.setxDesc(std::get(tensors))
+ .setdwDesc(std::get(tensors))
+ .setdyDesc(std::get(tensors))
+ .setcDesc(convDesc)
+ .setAlpha(alpha)
+ .setBeta(beta);
+ }
+ auto conv_op = conv_op_builder.build();
+ DEBUG_CUDNN_MSG(log_buf, conv_op.describe());
+
+ // Create an Operation Graph. In this case it is convolution add bias activation
+ std::array ops = {&conv_op};
+
+ auto opGraph = cudnn_frontend::OperationGraphBuilder()
+ .setHandle(handle_)
+ .setOperationGraph(ops.size(), ops.data())
+ .build();
+
+ // Create string encoding for plan caching
+ auto cache_string = getConvFusionString(x_dim_padded, pad, convstride, dilation, w_dim_padded, dataType, opGraph.getTag());
+ DEBUG_CUDNN_MSG(log_buf, "[convstring] " << cache_string);
+
+ auto& plan = getOrCreatePlan(handle_, log_buf, opGraph, cache_string);
+ DEBUG_CUDNN_MSG(log_buf, "Plan tag: " << plan.getTag());
+
+ auto workspace_size = plan.getWorkspaceSize();
+ DEBUG_CUDNN_MSG(log_buf, plan.describe() << " requires workspace " << workspace_size);
+
+ void* workspace_ptr = nullptr;
+ auto workspace_tensor = at::empty({(workspace_size+3)/4}, at::TensorOptions(at::kCUDA).dtype(at::kFloat));
+ if (workspace_size > 0) {
+ workspace_ptr = workspace_tensor.data_ptr();
+ }
+ void* data_ptrs[] = {devPtrX, devPtrY, devPtrW};
+ int64_t uids[] = {'x', 'y', 'w'};
+ auto variantPack = cudnn_frontend::VariantPackBuilder()
+ .setWorkspacePointer(workspace_ptr)
+ .setDataPointers(3, data_ptrs)
+ .setUids(3, uids)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, "variantPack " << variantPack.describe());
+ cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc());
+ checkCudnnErr(status);
+ cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, "Plan execute error", status);
+ } catch (cudnn_frontend::cudnnException e) {
+ std::cout << log_buf.str() << "[ERROR] Exception " << e.what() << std::endl;
+ }
+}
+
+void
+run_dconv_add(int64_t* x_dim_padded,
+ int64_t* pad,
+ int64_t* convstride,
+ int64_t* dilation,
+ int64_t* w_dim_padded,
+ int64_t* y_dim_padded,
+ cudnnDataType_t dataType,
+ at::Half* devPtrX,
+ at::Half* devPtrW,
+ at::Half* devPtrY,
+ at::Half* devPtrR) {
+ cudnnHandle_t handle_ = torch::native::getCudnnHandle();
+ std::stringstream log_buf;
+ try {
+ int convDim = 2;
+
+ // Creates the necessary tensor descriptors
+ dconv_descriptors tensors = create_dconv_descriptors(
+ x_dim_padded, pad, convstride, dilation, w_dim_padded, y_dim_padded, dataType);
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+
+ // Define the convolution problem
+ auto convDesc = cudnn_frontend::ConvDescBuilder()
+ .setDataType(CUDNN_DATA_FLOAT)
+ .setMathMode(CUDNN_CROSS_CORRELATION)
+ .setNDims(convDim)
+ .setStrides(convDim, convstride)
+ .setPrePadding(convDim, pad)
+ .setPostPadding(convDim, pad)
+ .setDilation(convDim, dilation)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, convDesc.describe());
+
+ // Define the add backward operation
+ auto addDesc = cudnn_frontend::PointWiseDescBuilder()
+ .setMode(CUDNN_POINTWISE_ADD)
+ .setMathPrecision(CUDNN_DATA_FLOAT)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, addDesc.describe());
+
+ float alpha = 1.0f;
+ float beta = 0.0f;
+
+ // Create a convolution Node
+ auto conv_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_DATA_DESCRIPTOR)
+ .setdxDesc(std::get(tensors))
+ .setwDesc(std::get(tensors))
+ .setdyDesc(std::get(tensors))
+ .setcDesc(convDesc)
+ .setAlpha(alpha)
+ .setBeta(beta)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, conv_op.describe());
+
+ // TODO: do we need getOutputTensor(), and what it returns in backward case?
+ // Create add Node.
+ auto add_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)
+ .setxDesc(std::get(tensors))
+ .setbDesc(std::get(tensors))
+ .setyDesc(std::get(tensors))
+ .setpwDesc(addDesc)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, add_op.describe());
+
+ // Create an Operation Graph. In this case it is convolution add bias activation
+ std::array ops = {&conv_op, &add_op};
+
+ auto opGraph = cudnn_frontend::OperationGraphBuilder()
+ .setHandle(handle_)
+ .setOperationGraph(ops.size(), ops.data())
+ .build();
+
+ // Create string encoding for plan caching
+ auto cache_string = getConvFusionString(x_dim_padded, pad, convstride, dilation, w_dim_padded, dataType, opGraph.getTag());
+ DEBUG_CUDNN_MSG(log_buf, "[convstring] " << cache_string);
+
+ auto& plan = getOrCreatePlan(handle_, log_buf, opGraph, cache_string);
+ DEBUG_CUDNN_MSG(log_buf, "Plan tag: " << plan.getTag());
+
+ auto workspace_size = plan.getWorkspaceSize();
+ DEBUG_CUDNN_MSG(log_buf, plan.describe() << " requires workspace " << workspace_size);
+
+ void* workspace_ptr = nullptr;
+ auto workspace_tensor = at::empty({(workspace_size+3)/4}, at::TensorOptions(at::kCUDA).dtype(at::kFloat));
+ if (workspace_size > 0) {
+ workspace_ptr = workspace_tensor.data_ptr();
+ }
+ void* data_ptrs[] = {devPtrX, devPtrY, devPtrW, devPtrR};
+ int64_t uids[] = {'x', 'y', 'w', 'r'};
+ auto variantPack = cudnn_frontend::VariantPackBuilder()
+ .setWorkspacePointer(workspace_ptr)
+ .setDataPointers(4, data_ptrs)
+ .setUids(4, uids)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, "variantPack " << variantPack.describe());
+ cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc());
+ checkCudnnErr(status);
+ cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, "Plan execute error", status);
+ } catch (cudnn_frontend::cudnnException e) {
+ std::cout << log_buf.str() << "[ERROR] Exception " << e.what() << std::endl;
+ }
+}
+
+
+// inputs contains x,w,z,b,(i)
+std::vector bottleneck_forward(bool explicit_nhwc, int stride_1X1, std::vector inputs) {
+
+ std::cout << std::fixed;
+ // create output vector
+ std::vector outputs;
+ auto output_format = explicit_nhwc ? at::MemoryFormat::Contiguous : at::MemoryFormat::ChannelsLast;
+
+ // setup dimensions
+ int64_t dimA[] = {0, 0, 0, 0};
+ int64_t filterdimA1[] = {0, 0, 0, 0};
+ int64_t filterdimA2[] = {0, 0, 0, 0};
+ int64_t filterdimA3[] = {0, 0, 0, 0};
+ int64_t filterdimA4[] = {0, 0, 0, 0};
+
+ // All dim calculation after this order of n,c,h,w
+ int axis[] {0,1,2,3};
+ if (explicit_nhwc) {
+ axis[0] = 0;
+ axis[1] = 3;
+ axis[2] = 1;
+ axis[3] = 2;
+ }
+ for (int dim=0;dim<4;dim++) {
+ dimA[dim] = inputs[0].size(axis[dim]);
+ filterdimA1[dim] = inputs[1].size(axis[dim]);
+ filterdimA2[dim] = inputs[2].size(axis[dim]);
+ filterdimA3[dim] = inputs[3].size(axis[dim]);
+ }
+ if (stride_1X1 != 1 || filterdimA3[0] != dimA[1]) {
+ for (int dim=0;dim<4;dim++) {
+ filterdimA4[dim] = inputs[10].size(axis[dim]);
+ }
+ }
+
+ // output dim in n,c,h,w used by backend
+ int64_t outdimA1[] = {0, 0, 0, 0}; // Computed Below
+ int64_t outdimA2[] = {0, 0, 0, 0}; // Computed Below
+ int64_t outdimA3[] = {0, 0, 0, 0}; // Computed Below
+
+ // use these fixed value for test run
+ int64_t padA[] = {0, 0};
+ int64_t padA1[] = {1, 1};
+ int64_t dilationA[] = {1, 1};
+ int64_t convstrideA[] = {1, 1};
+ int64_t convstride1X1[] = {stride_1X1, stride_1X1};
+
+ // compute output from pad/stride/dilation
+ outdimA1[0] = dimA[0];
+ outdimA1[1] = filterdimA1[0];
+ for (int dim = 0; dim < 2; dim++) {
+ outdimA1[dim + 2] = getFwdConvOutputDim(dimA[dim + 2], padA[dim], filterdimA1[dim + 2], convstride1X1[dim], dilationA[dim]);
+ }
+
+ outdimA2[0] = outdimA1[0];
+ outdimA2[1] = filterdimA2[0];
+ for (int dim = 0; dim < 2; dim++) {
+ outdimA2[dim + 2] = getFwdConvOutputDim(outdimA1[dim + 2], padA1[dim], filterdimA2[dim + 2], convstrideA[dim], dilationA[dim]);
+ }
+
+ outdimA3[0] = outdimA2[0];
+ outdimA3[1] = filterdimA3[0];
+ for (int dim = 0; dim < 2; dim++) {
+ outdimA3[dim + 2] = getFwdConvOutputDim(outdimA2[dim + 2], padA[dim], filterdimA3[dim + 2], convstrideA[dim], dilationA[dim]);
+ }
+
+ // Create output tensor in the correct shape in pytorch's view
+ int64_t outdim1[] = {0, 0, 0, 0};
+ int64_t outdim2[] = {0, 0, 0, 0};
+ int64_t outdim3[] = {0, 0, 0, 0};
+ if (explicit_nhwc) {
+ axis[0] = 0;
+ axis[1] = 2;
+ axis[2] = 3;
+ axis[3] = 1;
+ }
+ for (int dim=0;dim<4;dim++) {
+ outdim1[dim] = outdimA1[axis[dim]];
+ outdim2[dim] = outdimA2[axis[dim]];
+ outdim3[dim] = outdimA3[axis[dim]];
+ }
+
+ // run
+ at::Half* x = inputs[0].data_ptr();
+ at::Half* w = inputs[1].data_ptr();
+ at::Half* z = inputs[4].data_ptr();
+ at::Half* b = inputs[7].data_ptr();
+ auto out1 = at::empty(outdim1, inputs[0].type(), output_format);
+ at::Half* y1 = out1.data_ptr();
+
+ run_conv_scale_bias_add_activation(dimA,
+ padA,
+ convstride1X1,
+ dilationA,
+ filterdimA1,
+ outdimA1,
+ CUDNN_DATA_HALF,
+ x,
+ w,
+ y1,
+ z,
+ b,
+ nullptr);
+
+ DEBUG_MSG("[DEBUG] new relu1 : " << out1.to(at::kFloat).sum().item());
+
+ w = inputs[2].data_ptr();
+ z = inputs[5].data_ptr();
+ b = inputs[8].data_ptr();
+ auto out2 = at::empty(outdim2, inputs[0].type(), output_format);
+ at::Half* y2 = out2.data_ptr();
+
+ run_conv_scale_bias_add_activation(outdimA1,
+ padA1,
+ convstrideA,
+ dilationA,
+ filterdimA2,
+ outdimA2,
+ CUDNN_DATA_HALF,
+ y1,
+ w,
+ y2,
+ z,
+ b,
+ nullptr);
+ DEBUG_MSG("[DEBUG] new relu2 : " << out2.to(at::kFloat).sum().item());
+
+ // create output of conv3
+ auto out3 = at::empty(outdim3, inputs[0].type(), output_format);
+ at::Half* y3 = out3.data_ptr();
+
+ // create output of conv4 that may exist
+ auto identity = at::empty_like(out3);
+ at::Half* yi = identity.data_ptr();
+
+ if (stride_1X1 != 1 || filterdimA3[0] != dimA[1]){
+
+ w = inputs[10].data_ptr();
+ z = inputs[11].data_ptr();
+ b = inputs[12].data_ptr();
+ run_conv_scale_bias(dimA,
+ padA,
+ convstride1X1,
+ dilationA,
+ filterdimA4,
+ outdimA3,
+ CUDNN_DATA_HALF,
+ x,
+ w,
+ yi,
+ z,
+ b);
+ DEBUG_MSG("[DEBUG] new downsample : " << identity.to(at::kFloat).sum().item());
+ }
+ else {
+ yi = x;
+ }
+
+ w = inputs[3].data_ptr();
+ z = inputs[6].data_ptr();
+ b = inputs[9].data_ptr();
+
+ run_conv_scale_bias_add_activation(outdimA2,
+ padA,
+ convstrideA,
+ dilationA,
+ filterdimA3,
+ outdimA3,
+ CUDNN_DATA_HALF,
+ y2,
+ w,
+ y3,
+ z,
+ b,
+ yi);
+ DEBUG_MSG("[DEBUG] new relu3 : " << out3.to(at::kFloat).sum().item());
+
+ outputs.push_back(out1);
+ outputs.push_back(out2);
+ outputs.push_back(out3);
+
+ return outputs;
+}
+
+std::vector bottleneck_backward(bool explicit_nhwc, int stride_1X1, std::vector inputs) {
+
+ bool requires_grad = inputs[0].requires_grad();
+
+ std::cout << std::fixed;
+ // create output vector
+ std::vector outputs;
+ auto output_format = explicit_nhwc ? at::MemoryFormat::Contiguous : at::MemoryFormat::ChannelsLast;
+
+ // setup dimensions
+ int64_t dimA[] = {0, 0, 0, 0};
+ int64_t filterdimA1[] = {0, 0, 0, 0};
+ int64_t filterdimA2[] = {0, 0, 0, 0};
+ int64_t filterdimA3[] = {0, 0, 0, 0};
+ int64_t filterdimA4[] = {0, 0, 0, 0};
+
+ // All dim calculation after this order of n,c,h,w
+ int axis[] {0,1,2,3};
+ if (explicit_nhwc) {
+ axis[0] = 0;
+ axis[1] = 3;
+ axis[2] = 1;
+ axis[3] = 2;
+ }
+ for (int dim=0;dim<4;dim++) {
+ dimA[dim] = inputs[0].size(axis[dim]);
+ filterdimA1[dim] = inputs[1].size(axis[dim]);
+ filterdimA2[dim] = inputs[2].size(axis[dim]);
+ filterdimA3[dim] = inputs[3].size(axis[dim]);
+ }
+ if (stride_1X1 != 1 || filterdimA3[0] != dimA[1]) {
+ for (int dim=0;dim<4;dim++) {
+ filterdimA4[dim] = inputs[14].size(axis[dim]);
+ }
+ }
+
+ // output dim in n,c,h,w used by backend
+ int64_t outdimA1[] = {0, 0, 0, 0}; // Computed Below
+ int64_t outdimA2[] = {0, 0, 0, 0}; // Computed Below
+ int64_t outdimA3[] = {0, 0, 0, 0}; // Computed Below
+
+ // use these fixed value for test run
+ int64_t padA[] = {0, 0};
+ int64_t padA1[] = {1, 1};
+ int64_t dilationA[] = {1, 1};
+ int64_t convstrideA[] = {1, 1};
+ int64_t convstride1X1[] = {stride_1X1, stride_1X1};
+
+ // compute output from pad/stride/dilation
+ outdimA1[0] = dimA[0];
+ outdimA1[1] = filterdimA1[0];
+ for (int dim = 0; dim < 2; dim++) {
+ outdimA1[dim + 2] = getFwdConvOutputDim(dimA[dim + 2], padA[dim], filterdimA1[dim + 2], convstride1X1[dim], dilationA[dim]);
+ }
+
+ outdimA2[0] = outdimA1[0];
+ outdimA2[1] = filterdimA2[0];
+ for (int dim = 0; dim < 2; dim++) {
+ outdimA2[dim + 2] = getFwdConvOutputDim(outdimA1[dim + 2], padA1[dim], filterdimA2[dim + 2], convstrideA[dim], dilationA[dim]);
+ }
+
+ outdimA3[0] = outdimA2[0];
+ outdimA3[1] = filterdimA3[0];
+ for (int dim = 0; dim < 2; dim++) {
+ outdimA3[dim + 2] = getFwdConvOutputDim(outdimA2[dim + 2], padA[dim], filterdimA3[dim + 2], convstrideA[dim], dilationA[dim]);
+ }
+
+ // Create output tensor in the correct shape in pytorch's view
+ int64_t outdim1[] = {0, 0, 0, 0};
+ int64_t outdim2[] = {0, 0, 0, 0};
+ int64_t outdim3[] = {0, 0, 0, 0};
+ if (explicit_nhwc) {
+ axis[0] = 0;
+ axis[1] = 2;
+ axis[2] = 3;
+ axis[3] = 1;
+ }
+ for (int dim=0;dim<4;dim++) {
+ outdim1[dim] = outdimA1[axis[dim]];
+ outdim2[dim] = outdimA2[axis[dim]];
+ outdim3[dim] = outdimA3[axis[dim]];
+ }
+
+ // dconv3+drelu2+dscale2
+ at::Half* conv_in = inputs[13].data_ptr();
+ at::Half* dy3 = inputs[10].data_ptr();
+
+ DEBUG_MSG("[DEBUG] new dconv3 : " << inputs[10].to(at::kFloat).sum().item());
+
+ // wgrad
+ auto wgrad3 = at::empty_like(inputs[3]);
+ at::Half* dw3 = wgrad3.data_ptr();
+ run_dconv(outdimA2,
+ padA,
+ convstrideA,
+ dilationA,
+ filterdimA3,
+ outdimA3,
+ CUDNN_DATA_HALF,
+ conv_in,
+ dw3,
+ dy3,
+ CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_FILTER_DESCRIPTOR);
+
+ // dgrad
+ auto grad_out2 = at::empty(outdim2, inputs[0].type(), output_format);
+ at::Half* dy2 = grad_out2.data_ptr();
+ at::Half* w = inputs[3].data_ptr();
+ at::Half* z = inputs[5].data_ptr();
+
+ at::Half* relu2 = inputs[13].data_ptr();
+
+ run_dconv_drelu_dscale(outdimA2,
+ padA,
+ convstrideA,
+ dilationA,
+ filterdimA3,
+ outdimA3,
+ CUDNN_DATA_HALF,
+ dy2,
+ w,
+ dy3,
+ z,
+ relu2);
+
+ DEBUG_MSG("[DEBUG] new dconv2 : " << grad_out2.to(at::kFloat).sum().item());
+
+ // dconv2+drelu1+dscale1
+ conv_in = inputs[12].data_ptr();
+
+ // wgrad
+ auto wgrad2 = at::empty_like(inputs[2]);
+ at::Half* dw2 = wgrad2.data_ptr();
+ run_dconv(outdimA1,
+ padA1,
+ convstrideA,
+ dilationA,
+ filterdimA2,
+ outdimA2,
+ CUDNN_DATA_HALF,
+ conv_in,
+ dw2,
+ dy2,
+ CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_FILTER_DESCRIPTOR);
+
+ // dgrad
+ auto grad_out1 = at::empty(outdim1, inputs[0].type(), output_format);
+ at::Half* dy1 = grad_out1.data_ptr();
+ w = inputs[2].data_ptr();
+ z = inputs[4].data_ptr();
+
+ at::Half* relu1 = inputs[12].data_ptr();
+ // fused dgrad
+ run_dconv_drelu_dscale(outdimA1,
+ padA1,
+ convstrideA,
+ dilationA,
+ filterdimA2,
+ outdimA2,
+ CUDNN_DATA_HALF,
+ dy1,
+ w,
+ dy2,
+ z,
+ relu1);
+
+/*
+ // backward strided conv cannot be fused
+ // if stride == 1 but channel changes, we can fuse here
+ if (stride_1X1 != 1){
+ // dgrad
+ run_dconv(outdimA1,
+ padA1,
+ convstride1X1,
+ dilationA,
+ filterdimA2,
+ outdimA2,
+ CUDNN_DATA_HALF,
+ dy1,
+ w,
+ dy2,
+ CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_DATA_DESCRIPTOR);
+
+ // mul fused mask
+ grad_out1.mul_(inputs[15]);
+ }
+ else {
+ at::Half* relu1 = inputs[12].data_ptr();
+ // fused dgrad
+ run_dconv_drelu_dscale(outdimA1,
+ padA1,
+ convstride1X1,
+ dilationA,
+ filterdimA2,
+ outdimA2,
+ CUDNN_DATA_HALF,
+ dy1,
+ w,
+ dy2,
+ z,
+ relu1);
+ }
+*/
+ DEBUG_MSG("[DEBUG] new dconv1 : " << grad_out1.to(at::kFloat).sum().item());
+
+ // create grads of conv4 that may exist
+ auto grad_x_conv4 = at::empty_like(inputs[0]);
+ at::Half* dx_conv4 = grad_x_conv4.data_ptr();
+ at::Tensor wgrad4;
+
+ // x used for dconv1 and dconv4 wgrad
+ at::Half* x = inputs[0].data_ptr();
+
+ if (stride_1X1 != 1 || filterdimA3[0] != dimA[1]){
+ w = inputs[14].data_ptr();
+ at::Half* dy_conv4 = inputs[11].data_ptr();
+ if (requires_grad) {
+ run_dconv(dimA,
+ padA,
+ convstride1X1,
+ dilationA,
+ filterdimA4,
+ outdimA3,
+ CUDNN_DATA_HALF,
+ dx_conv4,
+ w,
+ dy_conv4,
+ CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_DATA_DESCRIPTOR);
+ // we don't print here since we can't hook out this grad in pytorch alone to compare, due to addition with dx
+ // DEBUG_MSG("[DEBUG] new dx_identity : " << grad_x_conv4.to(at::kFloat).sum().item());
+ }
+ // wgrad
+ wgrad4 = at::empty_like(inputs[14]);
+ at::Half* dw4 = wgrad4.data_ptr();
+ run_dconv(dimA,
+ padA,
+ convstride1X1,
+ dilationA,
+ filterdimA4,
+ outdimA3,
+ CUDNN_DATA_HALF,
+ x,
+ dw4,
+ dy_conv4,
+ CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_FILTER_DESCRIPTOR);
+ }
+ else {
+ // if there is no downsample, dx_conv4 is fork of drelu3
+ dx_conv4 = inputs[11].data_ptr();
+ }
+
+ // dconv1+add
+ // wgrad
+ auto wgrad1 = at::empty_like(inputs[1]);
+ at::Half* dw1 = wgrad1.data_ptr();
+ run_dconv(dimA,
+ padA,
+ convstride1X1,
+ dilationA,
+ filterdimA1,
+ outdimA1,
+ CUDNN_DATA_HALF,
+ x,
+ dw1,
+ dy1,
+ CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_FILTER_DESCRIPTOR);
+
+ // dgrad
+ w = inputs[1].data_ptr();
+ auto grad_x = at::empty_like(inputs[0]);
+ at::Half* dx = grad_x.data_ptr();
+
+ // backward strided conv cannot be fused
+ // if stride == 1 but channel changes, we can fuse here
+ if (requires_grad){
+ if (stride_1X1 != 1){
+ run_dconv(dimA,
+ padA,
+ convstride1X1,
+ dilationA,
+ filterdimA1,
+ outdimA1,
+ CUDNN_DATA_HALF,
+ dx,
+ w,
+ dy1,
+ CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_DATA_DESCRIPTOR);
+ // add 2 together
+ grad_x.add_(grad_x_conv4);
+ }
+ else {
+ run_dconv_add(dimA,
+ padA,
+ convstride1X1,
+ dilationA,
+ filterdimA1,
+ outdimA1,
+ CUDNN_DATA_HALF,
+ dx,
+ w,
+ dy1,
+ dx_conv4);
+ }
+ }
+
+ DEBUG_MSG("[DEBUG] new dx : " << grad_x.to(at::kFloat).sum().item());
+ DEBUG_MSG("[DEBUG] new wgrad1 : " << wgrad1.to(at::kFloat).sum().item());
+ DEBUG_MSG("[DEBUG] new wgrad2 : " << wgrad2.to(at::kFloat).sum().item());
+ DEBUG_MSG("[DEBUG] new wgrad3 : " << wgrad3.to(at::kFloat).sum().item());
+ outputs.push_back(grad_x);
+ outputs.push_back(wgrad1);
+ outputs.push_back(wgrad2);
+ outputs.push_back(wgrad3);
+
+ if (stride_1X1 != 1 || filterdimA3[0] != dimA[1]) {
+ DEBUG_MSG("[DEBUG] new wgrad4 : " << wgrad4.to(at::kFloat).sum().item());
+ outputs.push_back(wgrad4);
+ }
+
+ return outputs;
+}
+
+namespace {
+
+enum {
+ X_TENSOR,
+ Y_TENSOR,
+ W_TENSOR,
+ Z_TENSOR,
+ B_TENSOR,
+ AFTERADD_TENSOR,
+ AFTERBIAS_TENSOR,
+ AFTERCONV_TENSOR,
+ OPTIONAL,
+ AFTEROPT_TENSOR,
+ AFTERACT_TENSOR,
+ GEN_INDEX_TENSOR,
+ MASK_TOP_TENSOR,
+ MASK_BOTTOM_TENSOR,
+ MASK_TENSOR,
+ THRESHOLD_TOP_TENSOR,
+ THRESHOLD_BOTTOM_TENSOR,
+};
+
+using masked_convbias_descriptors = std::tuple;
+
+masked_convbias_descriptors
+create_conv_bias_add_act_mask_descriptors(int64_t* x_dim_padded,
+ int64_t* padA,
+ int64_t* convstrideA,
+ int64_t* dilationA,
+ int64_t* w_dim_padded,
+ int64_t* y_dim_padded,
+ int64_t* threshold_dim,
+ cudnnDataType_t dataType) {
+ const int convDim = 2;
+
+ int64_t b_dim_padded[4];
+ b_dim_padded[0] = 1;
+ b_dim_padded[1] = y_dim_padded[1];
+ b_dim_padded[2] = 1;
+ b_dim_padded[3] = 1;
+
+ int64_t x_stride_padded[4];
+ int64_t y_stride_padded[4];
+ int64_t w_stride_padded[4];
+ int64_t b_stride_padded[4];
+ int64_t threshold_stride[4];
+
+ generateStrides(w_dim_padded, w_stride_padded, 4, CUDNN_TENSOR_NHWC);
+ generateStrides(x_dim_padded, x_stride_padded, 4, CUDNN_TENSOR_NHWC);
+ generateStrides(y_dim_padded, y_stride_padded, 4, CUDNN_TENSOR_NHWC);
+ generateStrides(b_dim_padded, b_stride_padded, 4, CUDNN_TENSOR_NHWC);
+ generateStrides(threshold_dim, threshold_stride, 4, CUDNN_TENSOR_NHWC);
+
+ return masked_convbias_descriptors(cudnn_frontend::TensorBuilder()
+ .setDim(4, x_dim_padded)
+ .setStrides(4, x_stride_padded)
+ .setId('x')
+ .setAlignment(16)
+ .setDataType(dataType)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, y_dim_padded)
+ .setStrides(4, y_stride_padded)
+ .setId('y')
+ .setAlignment(16)
+ .setDataType(dataType)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, w_dim_padded)
+ .setStrides(4, w_stride_padded)
+ .setId('w')
+ .setAlignment(16)
+ .setDataType(dataType)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, b_dim_padded)
+ .setStrides(4, b_stride_padded)
+ .setId('z')
+ .setAlignment(16)
+ .setDataType(dataType)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, b_dim_padded)
+ .setStrides(4, b_stride_padded)
+ .setId('b')
+ .setAlignment(16)
+ .setDataType(dataType)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, y_dim_padded)
+ .setStrides(4, y_stride_padded)
+ .setVirtual()
+ .setId('A') // after add
+ .setAlignment(16)
+ .setDataType(CUDNN_DATA_FLOAT)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, y_dim_padded)
+ .setStrides(4, y_stride_padded)
+ .setVirtual()
+ .setId('B') // after bias
+ .setAlignment(16)
+ .setDataType(CUDNN_DATA_FLOAT)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, y_dim_padded)
+ .setStrides(4, y_stride_padded)
+ .setId('C') // after conv
+ .setAlignment(16)
+ .setVirtual()
+ .setDataType(CUDNN_DATA_FLOAT)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, y_dim_padded)
+ .setStrides(4, y_stride_padded)
+ .setId('i')
+ .setAlignment(16)
+ .setDataType(dataType)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, y_dim_padded)
+ .setStrides(4, y_stride_padded)
+ .setId('D') // after optional add
+ .setAlignment(16)
+ .setVirtual()
+ .setDataType(CUDNN_DATA_FLOAT)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, y_dim_padded)
+ .setStrides(4, y_stride_padded)
+ .setId('E') // after act for masked
+ .setAlignment(16)
+ .setVirtual()
+ .setDataType(CUDNN_DATA_FLOAT)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, y_dim_padded)
+ .setStrides(4, y_stride_padded)
+ .setId('I') // output of the gen index operation
+ .setAlignment(16)
+ .setVirtual()
+ .setDataType(CUDNN_DATA_INT32)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, y_dim_padded)
+ .setStrides(4, y_stride_padded)
+ .setId('m') // top half of the mask created after the less than
+ .setAlignment(16)
+ .setVirtual()
+ .setDataType(CUDNN_DATA_BOOLEAN)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, y_dim_padded)
+ .setStrides(4, y_stride_padded)
+ .setId('n') // bottom half of the mask
+ .setAlignment(16)
+ .setVirtual()
+ .setDataType(CUDNN_DATA_BOOLEAN)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, y_dim_padded)
+ .setStrides(4, y_stride_padded)
+ .setId('M') // OR of the top and bottom masks
+ .setAlignment(16)
+ .setVirtual()
+ .setDataType(CUDNN_DATA_BOOLEAN)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, threshold_dim)
+ .setStrides(4, threshold_stride)
+ .setId('t') // threshold for creating the top mask
+ .setAlignment(16)
+ .setDataType(CUDNN_DATA_INT32)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, threshold_dim)
+ .setStrides(4, threshold_stride)
+ .setId('u') // threshold for creating the bottom mask
+ .setAlignment(16)
+ .setDataType(CUDNN_DATA_INT32)
+ .build());
+}
+
+// tensor descriptors used for dgrad
+enum {
+ X_OR_DX_TENSOR,
+ DY_TENSOR,
+ W_OR_DW_TENSOR,
+ SCALE_TENSOR,
+ RELU_TENSOR,
+ AFTER_DCONV_TENSOR,
+ AFTER_DRELU_TENSOR,
+ DGRAD_INPUT_TENSOR,
+ DGRAD_OPTIONAL_TENSOR,
+ DGRAD_GEN_INDEX_TENSOR,
+ DGRAD_MASK_TOP_TENSOR,
+ DGRAD_MASK_BOTTOM_TENSOR,
+ DGRAD_MASK_TENSOR,
+ DGRAD_THRESHOLD_TOP_TENSOR,
+ DGRAD_THRESHOLD_BOTTOM_TENSOR,
+};
+
+using dconv_add_descriptors = std::tuple;
+
+dconv_add_descriptors
+create_dconv_add_descriptors(int64_t* x_dim_padded,
+ int64_t* padA,
+ int64_t* convstrideA,
+ int64_t* dilationA,
+ int64_t* w_dim_padded,
+ int64_t* y_dim_padded,
+ cudnnDataType_t dataType) {
+ const int convDim = 2;
+
+ int64_t b_dim_padded[4];
+ b_dim_padded[0] = 1;
+ b_dim_padded[1] = x_dim_padded[1];
+ b_dim_padded[2] = 1;
+ b_dim_padded[3] = 1;
+
+ int64_t x_stride_padded[4];
+ int64_t y_stride_padded[4];
+ int64_t w_stride_padded[4];
+ int64_t b_stride_padded[4];
+
+ generateStrides(w_dim_padded, w_stride_padded, 4, CUDNN_TENSOR_NHWC);
+ generateStrides(x_dim_padded, x_stride_padded, 4, CUDNN_TENSOR_NHWC);
+ generateStrides(y_dim_padded, y_stride_padded, 4, CUDNN_TENSOR_NHWC);
+ generateStrides(b_dim_padded, b_stride_padded, 4, CUDNN_TENSOR_NHWC);
+
+ return dconv_add_descriptors(cudnn_frontend::TensorBuilder()
+ .setDim(4, x_dim_padded)
+ .setStrides(4, x_stride_padded)
+ .setId('x')
+ .setAlignment(16)
+ .setDataType(dataType)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, y_dim_padded)
+ .setStrides(4, y_stride_padded)
+ .setId('y')
+ .setAlignment(16)
+ .setDataType(dataType)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, w_dim_padded)
+ .setStrides(4, w_stride_padded)
+ .setId('w')
+ .setAlignment(16)
+ .setDataType(dataType)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, b_dim_padded)
+ .setStrides(4, b_stride_padded)
+ .setId('s')
+ .setAlignment(16)
+ .setDataType(dataType)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, x_dim_padded)
+ .setStrides(4, x_stride_padded)
+ .setId('r')
+ .setAlignment(16)
+ .setDataType(dataType)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, x_dim_padded)
+ .setStrides(4, x_stride_padded)
+ .setVirtual()
+ .setId('A') // after dconv
+ .setAlignment(16)
+ .setDataType(CUDNN_DATA_FLOAT)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, x_dim_padded)
+ .setStrides(4, x_stride_padded)
+ .setVirtual()
+ .setId('B') // after drelu
+ .setAlignment(16)
+ .setDataType(CUDNN_DATA_FLOAT)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, y_dim_padded)
+ .setStrides(4, y_stride_padded)
+ .setId('i')
+ .setAlignment(16)
+ .setDataType(dataType)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, y_dim_padded)
+ .setStrides(4, y_stride_padded)
+ .setId('D') // after optional add
+ .setAlignment(16)
+ .setVirtual()
+ .setDataType(CUDNN_DATA_FLOAT)
+ .build());
+}
+
+using dconv_mask_descriptors = std::tuple;
+
+dconv_mask_descriptors
+create_dconv_mask_descriptors(int64_t* x_dim_padded,
+ int64_t* padA,
+ int64_t* convstrideA,
+ int64_t* dilationA,
+ int64_t* w_dim_padded,
+ int64_t* y_dim_padded,
+ int64_t* threshold_dim,
+ cudnnDataType_t dataType) {
+ const int convDim = 2;
+
+ int64_t b_dim_padded[4];
+ b_dim_padded[0] = 1;
+ b_dim_padded[1] = x_dim_padded[1];
+ b_dim_padded[2] = 1;
+ b_dim_padded[3] = 1;
+
+ int64_t x_stride_padded[4];
+ int64_t y_stride_padded[4];
+ int64_t w_stride_padded[4];
+ int64_t b_stride_padded[4];
+ int64_t threshold_stride[4];
+
+ generateStrides(w_dim_padded, w_stride_padded, 4, CUDNN_TENSOR_NHWC);
+ generateStrides(x_dim_padded, x_stride_padded, 4, CUDNN_TENSOR_NHWC);
+ generateStrides(y_dim_padded, y_stride_padded, 4, CUDNN_TENSOR_NHWC);
+ generateStrides(b_dim_padded, b_stride_padded, 4, CUDNN_TENSOR_NHWC);
+ generateStrides(threshold_dim, threshold_stride, 4, CUDNN_TENSOR_NHWC);
+
+ return dconv_mask_descriptors(cudnn_frontend::TensorBuilder()
+ .setDim(4, x_dim_padded)
+ .setStrides(4, x_stride_padded)
+ .setId('x')
+ .setAlignment(16)
+ .setDataType(dataType)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, y_dim_padded)
+ .setStrides(4, y_stride_padded)
+ .setId('y')
+ .setAlignment(16)
+ .setDataType(dataType)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, w_dim_padded)
+ .setStrides(4, w_stride_padded)
+ .setId('w')
+ .setAlignment(16)
+ .setDataType(dataType)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, b_dim_padded)
+ .setStrides(4, b_stride_padded)
+ .setId('s')
+ .setAlignment(16)
+ .setDataType(dataType)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, x_dim_padded)
+ .setStrides(4, x_stride_padded)
+ .setId('r')
+ .setAlignment(16)
+ .setDataType(dataType)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, x_dim_padded)
+ .setStrides(4, x_stride_padded)
+ .setVirtual()
+ .setId('A') // after dconv
+ .setAlignment(16)
+ .setDataType(CUDNN_DATA_FLOAT)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, x_dim_padded)
+ .setStrides(4, x_stride_padded)
+ .setVirtual()
+ .setId('B') // after drelu
+ .setAlignment(16)
+ .setDataType(CUDNN_DATA_FLOAT)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, y_dim_padded)
+ .setStrides(4, y_stride_padded)
+ .setId('i')
+ .setAlignment(16)
+ .setDataType(dataType)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, y_dim_padded)
+ .setStrides(4, y_stride_padded)
+ .setId('D') // after optional add
+ .setAlignment(16)
+ .setVirtual()
+ .setDataType(CUDNN_DATA_FLOAT)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, y_dim_padded)
+ .setStrides(4, y_stride_padded)
+ .setId('I') // output of the gen index operation
+ .setAlignment(16)
+ .setVirtual()
+ .setDataType(CUDNN_DATA_INT32)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, y_dim_padded)
+ .setStrides(4, y_stride_padded)
+ .setId('m') // top half of the mask created after the less than
+ .setAlignment(16)
+ .setVirtual()
+ .setDataType(CUDNN_DATA_BOOLEAN)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, y_dim_padded)
+ .setStrides(4, y_stride_padded)
+ .setId('n') // bottom half of the mask
+ .setAlignment(16)
+ .setVirtual()
+ .setDataType(CUDNN_DATA_BOOLEAN)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, y_dim_padded)
+ .setStrides(4, y_stride_padded)
+ .setId('M') // OR of the top and bottom masks
+ .setAlignment(16)
+ .setVirtual()
+ .setDataType(CUDNN_DATA_BOOLEAN)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, threshold_dim)
+ .setStrides(4, threshold_stride)
+ .setId('t') // threshold for creating the top mask
+ .setAlignment(16)
+ .setDataType(CUDNN_DATA_INT32)
+ .build(),
+ cudnn_frontend::TensorBuilder()
+ .setDim(4, threshold_dim)
+ .setStrides(4, threshold_stride)
+ .setId('u') // threshold for creating the bottom mask
+ .setAlignment(16)
+ .setDataType(CUDNN_DATA_INT32)
+ .build());
+}
+
+void
+run_conv_add_scale_bias_activation(int64_t* x_dim_padded,
+ int64_t* pad,
+ int64_t* convstride,
+ int64_t* dilation,
+ int64_t* w_dim_padded,
+ int64_t* y_dim_padded,
+ cudnnDataType_t dataType,
+ at::Half* devPtrX,
+ at::Half* devPtrW,
+ at::Half* devPtrY,
+ at::Half* devPtrZ,
+ at::Half* devPtrB,
+ at::Half* devPtrI) {
+ cudnnHandle_t handle_ = torch::native::getCudnnHandle();
+ std::stringstream log_buf;
+ try {
+ int convDim = 2;
+
+ // Creates the necessary tensor descriptors
+ common_convbias_descriptors tensors = create_conv_bias_add_act_descriptors(
+ x_dim_padded, pad, convstride, dilation, w_dim_padded, y_dim_padded, dataType);
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+
+ // Define the add operation
+ auto scaleDesc = cudnn_frontend::PointWiseDescBuilder()
+ .setMode(CUDNN_POINTWISE_MUL)
+ .setMathPrecision(CUDNN_DATA_FLOAT)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, scaleDesc.describe());
+
+ // Define the bias operation
+ auto biasDesc = cudnn_frontend::PointWiseDescBuilder()
+ .setMode(CUDNN_POINTWISE_ADD)
+ .setMathPrecision(CUDNN_DATA_FLOAT)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, biasDesc.describe());
+
+ // optional add
+ auto addDesc = cudnn_frontend::PointWiseDescBuilder()
+ .setMode(CUDNN_POINTWISE_ADD)
+ .setMathPrecision(CUDNN_DATA_FLOAT)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, addDesc.describe());
+
+ // Define the activation operation
+ auto actDesc = cudnn_frontend::PointWiseDescBuilder()
+ .setMode(CUDNN_POINTWISE_RELU_FWD)
+ .setMathPrecision(CUDNN_DATA_FLOAT)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, actDesc.describe());
+
+ // Define the convolution problem
+ auto convDesc = cudnn_frontend::ConvDescBuilder()
+ .setDataType(CUDNN_DATA_FLOAT)
+ .setMathMode(CUDNN_CROSS_CORRELATION)
+ .setNDims(convDim)
+ .setStrides(convDim, convstride)
+ .setPrePadding(convDim, pad)
+ .setPostPadding(convDim, pad)
+ .setDilation(convDim, dilation)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, convDesc.describe());
+
+ float alpha = 1.0f;
+ float beta = 0.0f;
+
+ // Create a convolution Node
+ auto conv_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_CONVOLUTION_FORWARD_DESCRIPTOR)
+ .setxDesc(std::get(tensors))
+ .setwDesc(std::get(tensors))
+ .setyDesc(std::get(tensors))
+ .setcDesc(convDesc)
+ .setAlpha(alpha)
+ .setBeta(beta)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, conv_op.describe());
+
+ // create an add node.
+ auto add_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)
+ .setxDesc(conv_op.getOutputTensor())
+ .setbDesc(std::get(tensors))
+ .setyDesc(std::get(tensors))
+ .setpwDesc(addDesc)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, add_op.describe());
+
+ // Create a Add Node with scaling parameters.
+ auto scale_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)
+ .setxDesc(add_op.getOutputTensor())
+ .setbDesc(std::get(tensors))
+ .setyDesc(std::get(tensors))
+ .setpwDesc(scaleDesc)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, scale_op.describe());
+
+ // Create a Bias Node.
+ auto bias_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)
+ .setxDesc(scale_op.getOutputTensor())
+ .setbDesc(std::get(tensors))
+ .setyDesc(std::get(tensors))
+ .setpwDesc(biasDesc)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, bias_op.describe());
+
+ // Create an Activation Node.
+ auto act_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)
+ .setxDesc(bias_op.getOutputTensor())
+ .setyDesc(std::get(tensors))
+ .setpwDesc(actDesc)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, act_op.describe());
+
+ // Create an Operation Graph. In this case it is convolution add bias activation
+ std::array ops = {&conv_op, &add_op, &scale_op, &bias_op, &act_op};
+
+ auto opGraph = cudnn_frontend::OperationGraphBuilder()
+ .setHandle(handle_)
+ .setOperationGraph(ops.size(), ops.data())
+ .build();
+
+ // Create string encoding for plan caching
+ auto cache_string = getConvFusionString(x_dim_padded, pad, convstride, dilation, w_dim_padded, dataType, opGraph.getTag());
+ DEBUG_CUDNN_MSG(log_buf, "[convstring] " << cache_string);
+
+ auto& plan = getOrCreatePlan(handle_, log_buf, opGraph, cache_string);
+ DEBUG_CUDNN_MSG(log_buf, "Plan tag: " << plan.getTag());
+
+ auto workspace_size = plan.getWorkspaceSize();
+ DEBUG_CUDNN_MSG(log_buf, plan.describe() << " requires workspace " << workspace_size);
+
+ void* workspace_ptr = nullptr;
+ auto workspace_tensor = at::empty({(workspace_size+3)/4}, at::TensorOptions(at::kCUDA).dtype(at::kFloat));
+ if (workspace_size > 0) {
+ workspace_ptr = workspace_tensor.data_ptr();
+ }
+ void* data_ptrs[] = {devPtrX, devPtrY, devPtrW, devPtrZ, devPtrB, devPtrI};
+ int64_t uids[] = {'x', 'y', 'w', 'z', 'b', 'i'};
+ auto variantPack = cudnn_frontend::VariantPackBuilder()
+ .setWorkspacePointer(workspace_ptr)
+ .setDataPointers(6, data_ptrs)
+ .setUids(6, uids)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, "variantPack " << variantPack.describe());
+ cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc());
+ checkCudnnErr(status);
+ cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, "Plan execute error", status);
+ } catch (cudnn_frontend::cudnnException e) {
+ std::cout << log_buf.str() << "[ERROR] Exception " << e.what() << std::endl;
+ }
+}
+
+void
+run_conv_scale_bias_add_activation_mask(int64_t* x_dim_padded,
+ int64_t* pad,
+ int64_t* convstride,
+ int64_t* dilation,
+ int64_t* w_dim_padded,
+ int64_t* y_dim_padded,
+ int64_t* threshold_dim,
+ cudnnDataType_t dataType,
+ at::Half* devPtrX,
+ at::Half* devPtrW,
+ at::Half* devPtrY,
+ at::Half* devPtrZ,
+ at::Half* devPtrB,
+ at::Half* devPtrI,
+ int* devPtrT,
+ int* devPtrU,
+ int axis) {
+ cudnnHandle_t handle_ = torch::native::getCudnnHandle();
+ std::stringstream log_buf;
+ try {
+ int convDim = 2;
+
+ // Creates the necessary tensor descriptors
+ masked_convbias_descriptors tensors = create_conv_bias_add_act_mask_descriptors(
+ x_dim_padded, pad, convstride, dilation, w_dim_padded, y_dim_padded, threshold_dim, dataType);
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+ DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());
+
+ // Define the add operation
+ auto scaleDesc = cudnn_frontend::PointWiseDescBuilder()
+ .setMode(CUDNN_POINTWISE_MUL)
+ .setMathPrecision(CUDNN_DATA_FLOAT)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, scaleDesc.describe());
+
+ // Define the bias operation
+ auto biasDesc = cudnn_frontend::PointWiseDescBuilder()
+ .setMode(CUDNN_POINTWISE_ADD)
+ .setMathPrecision(CUDNN_DATA_FLOAT)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, biasDesc.describe());
+
+ // optional add
+ auto addDesc = cudnn_frontend::PointWiseDescBuilder()
+ .setMode(CUDNN_POINTWISE_ADD)
+ .setMathPrecision(CUDNN_DATA_FLOAT)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, addDesc.describe());
+
+ // Define the activation operation
+ auto actDesc = cudnn_frontend::PointWiseDescBuilder()
+ .setMode(CUDNN_POINTWISE_RELU_FWD)
+ .setMathPrecision(CUDNN_DATA_FLOAT)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, actDesc.describe());
+
+ // Define the convolution problem
+ auto convDesc = cudnn_frontend::ConvDescBuilder()
+ .setDataType(CUDNN_DATA_FLOAT)
+ .setMathMode(CUDNN_CROSS_CORRELATION)
+ .setNDims(convDim)
+ .setStrides(convDim, convstride)
+ .setPrePadding(convDim, pad)
+ .setPostPadding(convDim, pad)
+ .setDilation(convDim, dilation)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, convDesc.describe());
+
+ // Define the genIndex descriptor
+ auto genIndexDesc = cudnn_frontend::PointWiseDescBuilder()
+ .setMode(CUDNN_POINTWISE_GEN_INDEX)
+ .setMathPrecision(CUDNN_DATA_FLOAT)
+ .setAxis(axis)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, genIndexDesc.describe());
+
+ // Define the lessThan descriptor
+ auto lessThanDesc = cudnn_frontend::PointWiseDescBuilder()
+ .setMode(CUDNN_POINTWISE_CMP_LT)
+ .setMathPrecision(CUDNN_DATA_FLOAT)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, lessThanDesc.describe());
+
+ // Define the greaterThan descriptor
+ auto greaterThanDesc = cudnn_frontend::PointWiseDescBuilder()
+ .setMode(CUDNN_POINTWISE_CMP_GT)
+ .setMathPrecision(CUDNN_DATA_FLOAT)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, greaterThanDesc.describe());
+
+ // Define the logical_or descriptor
+ auto logicalOrDesc = cudnn_frontend::PointWiseDescBuilder()
+ .setMode(CUDNN_POINTWISE_LOGICAL_OR)
+ .setMathPrecision(CUDNN_DATA_BOOLEAN)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, logicalOrDesc.describe());
+
+ // Define the binary_selection descriptor
+ auto selectionDesc = cudnn_frontend::PointWiseDescBuilder()
+ .setMode(CUDNN_POINTWISE_BINARY_SELECT)
+ .setMathPrecision(CUDNN_DATA_FLOAT)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, selectionDesc.describe());
+
+ float alpha = 1.0f;
+ float beta = 0.0f;
+
+ // Create a convolution Node
+ auto conv_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_CONVOLUTION_FORWARD_DESCRIPTOR)
+ .setxDesc(std::get(tensors))
+ .setwDesc(std::get(tensors))
+ .setyDesc(std::get(tensors))
+ .setcDesc(convDesc)
+ .setAlpha(alpha)
+ .setBeta(beta)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, conv_op.describe());
+
+ // Create a Add Node with scaling parameters.
+ auto scale_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)
+ .setxDesc(conv_op.getOutputTensor())
+ .setbDesc(std::get(tensors))
+ .setyDesc(std::get(tensors))
+ .setpwDesc(scaleDesc)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, scale_op.describe());
+
+ // Create a Bias Node.
+ auto bias_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)
+ .setxDesc(scale_op.getOutputTensor())
+ .setbDesc(std::get(tensors))
+ .setyDesc(std::get(tensors))
+ .setpwDesc(biasDesc)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, bias_op.describe());
+
+ // Create a optional add Node.
+ auto add_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)
+ .setxDesc(bias_op.getOutputTensor())
+ .setbDesc(std::get(tensors))
+ .setyDesc(std::get(tensors))
+ .setpwDesc(addDesc)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, add_op.describe());
+
+
+ // Create an Activation Node.
+ auto act_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)
+ .setxDesc(devPtrI ? add_op.getOutputTensor() : bias_op.getOutputTensor())
+ .setyDesc(std::get(tensors))
+ .setpwDesc(actDesc)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, act_op.describe());
+
+ // Create a Gen_Index Node.
+ auto genIndex_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)
+ .setxDesc(std::get(tensors))
+ .setyDesc(std::get(tensors))
+ .setpwDesc(genIndexDesc)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, genIndex_op.describe());
+
+ // Create a LessThan Node.
+ auto lessThan_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)
+ .setxDesc(std::get(tensors))
+ .setbDesc(std::get(tensors))
+ .setyDesc(std::get(tensors))
+ .setpwDesc(lessThanDesc)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, lessThan_op.describe());
+
+ // Create a GreaterThan Node.
+ auto greaterThan_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)
+ .setxDesc(std::get(tensors))
+ .setbDesc(std::get(tensors))
+ .setyDesc(std::get(tensors))
+ .setpwDesc(greaterThanDesc)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, greaterThan_op.describe());
+
+ // Create a LogicalOr Node.
+ auto logicalOr_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)
+ .setxDesc(std::get(tensors))
+ .setbDesc(std::get(tensors))
+ .setyDesc(std::get(tensors))
+ .setpwDesc(logicalOrDesc)
+ .build();
+ DEBUG_CUDNN_MSG(log_buf, logicalOr_op.describe());
+
+ // Create a Binary_Selection Node.
+ auto selection_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)
+ .setxDesc(std::get(tensors))
+ .setbDesc(std::get(tensors))
+ .settDesc(std::get(tensors))
+ .setyDesc(std::get 
- 
-
-
-
- 
-
-
-
- 
- 
- 
- 
- 
- 
- 
- 
- 
- 
- 
-