Tinh chỉnh OpenFlamingo với GPU NVIDIA H100

1. Giới thiệu Flamingo: Few-shot learning cho các mô hình ngôn ngữ thị giác 

Flamingo (tài liệu gốc: [https://arxiv.org/pdf/2204.14198]) là một nhánh của mô hình ngôn ngữ thị giác (VLMs) được xây dựng bởi nhóm các chuyên gia tại Google DeepMind để giải quyết các bài toán về Few-shot learning trong học máy đa phương thức. Mô hình này được xây dựng dựa trên ba tầng lớp quan trọng: 

• Kết hợp sức mạnh của các mô hình đã huấn luyện sẵn cho xử lý hình ảnh và xử lý ngôn ngữ. 

• Xử lý linh hoạt chuỗi dữ liệu gồm cả hình ảnh và văn bản xen kẽ. 

• Dễ dàng nhận hình ảnh hoặc video làm đầu vào. 

Sự linh hoạt này giúp Flamingo có thể được huấn luyện trên khối lượng dữ liệu web quy mô lớn, nơi hình ảnh và văn bản được kết hợp lẫn nhau. Đây chính là yếu tố quan trọng để Flamingo có khả năng học những nhiệm vụ mới chỉ với một vài ví dụ. 

Nhờ vậy, bằng cách sử dụng một vài ví dụ mẫu phù hợp với từng tác vụ, một mô hình Flamingo có thể đạt được hiệu năng hàng đầu trong nhiều tác vụ khác nhau, bao gồm trả lời câu hỏi dựa trên hình ảnh, tạo chú thích (caption) và trả lời câu hỏi trắc nghiệm. 

Cách tiếp cận few-shot này thường giúp Flamingo vượt trội hơn cả những mô hình đã được tinh chỉnh bằng lượng dữ liệu nhiều gấp hàng nghìn lần. 

2. Flamingo vận hành như thế nào?

Flamingo hoạt động thông qua một giao diện đa phương thức, xử lý kết hợp hình ảnh, video và văn bản để tạo ra phản hồi dưới dạng văn bản phù hợp. Thiết kế này giúp Flamingo dễ dàng thích ứng với nhiều tác vụ khác nhau, tương tự như các mô hình ngôn ngữ lớn (LLM), vốn sử dụng các gợi ý bằng văn bản để giải quyết nhiều thách thức liên quan đến ngôn ngữ. 

Kiến trúc mô hình  

OpenFlamingo kết hợp một bộ mã hóa thị giác đã được huấn luyện sẵn với một mô hình ngôn ngữ thông qua các lớp cross-attention. Kiến trúc mô hình được minh họa bên dưới. 

Kiến trúc này hoạt động thông qua hai nhánh chính: 

 a. Nhánh thị giác (bên trái)

Nhánh này chịu trách nhiệm xử lý dữ liệu thị giác (hình ảnh) và chuẩn bị cho mô hình ngôn ngữ. 

• Vision Encoder: Đây là một mô hình đã được huấn luyện sẵn (biểu thị bằng biểu tượng bông tuyết “đóng băng”), có nhiệm vụ trích xuất đặc trưng từ hình ảnh đầu vào. Một điểm quan trọng trong thiết kế là các trọng số của bộ mã hóa này được “đóng băng” và không thay đổi trong quá trình huấn luyện. 

• Perceiver Resampler: Đầu ra từ bộ mã hóa thị giác sau đó được đưa vào Perceiver Resampler. Thành phần này ánh xạ các đặc trưng thị giác có kích thước thay đổi thành một số lượng nhỏ, cố định các token đầu ra. Mô-đun này được huấn luyện từ đầu (biểu thị bằng màu tím), học cách tạo ra phiên bản rút gọn của dữ liệu hình ảnh. Trong Flamingo, số lượng token hình ảnh đầu ra được đặt là 5. 

b. Nhánh ngôn ngữ (bên phải)

Nhánh này xử lý văn bản và kết hợp nó với thông tin thị giác để tạo ra đầu ra cuối cùng. 

• Interleaved Input: Mô hình nhận một chuỗi đầu vào gồm văn bản xen kẽ với các ký hiệu giữ chỗ cho hình ảnh (<image>). 

• LM Blocks: Nền tảng của nhánh này là một mô hình ngôn ngữ lớn (Language Model – LM) đã được huấn luyện sẵn (ví dụ: Chinchilla). Tương tự như Bộ mã hóa thị giác, các khối LM này cũng được “đóng băng”, nghĩa là tận dụng vốn kiến thức ngôn ngữ khổng lồ mà không cần huấn luyện lại. 

• Gated XATTN-DENSE: Đây là điểm đổi mới then chốt giúp kết nối hai nhánh. Đây là những mô-đun mới, được huấn luyện từ đầu và chèn vào giữa các khối LM. Khi mô hình gặp một ký hiệu <image> trong dòng văn bản, lớp Gated XATTN-DENSE sẽ thực hiện phép cross-attention, dùng thông tin văn bản làm “truy vấn” để “nhìn vào” các token hình ảnh do Perceiver Resampler sinh ra. Thành phần “gated” là cơ chế kiểm soát mức độ thông tin hình ảnh được phép ảnh hưởng đến quá trình sinh ngôn ngữ, giúp kết hợp linh hoạt hai loại dữ liệu. 

Thiết lập tiêu chuẩn mới cho Few-Shot learning 

Flamingo đã được kiểm chứng nghiêm ngặt trên 16 tác vụ khác nhau và luôn đạt kết quả tốt hơn so với các mô hình few-shot learning trước đây, ngay cả khi mỗi tác vụ chỉ có bốn ví dụ. Trong một số trường hợp, Flamingo còn cho thấy hiệu năng tốt hơn cả những phương pháp cần tinh chỉnh trên quy mô lớn và sử dụng tập dữ liệu khổng lồ, cho thấy khả năng khái quát hóa rất hiệu quả. 

Bằng cách giảm thiểu nhu cầu gắn nhãn dữ liệu quy mô lớn và huấn luyện lại riêng cho từng tác vụ, Flamingo đánh dấu một bước tiến quan trọng về hiệu quả của các mô hình ngôn ngữ thị giác. Khả năng học nhanh từ số lượng ví dụ hạn chế đưa AI tiến gần hơn đến sự thích ứng tự nhiên như con người, mở ra nhiều ứng dụng thực tiễn với độ chính xác cao và dễ dàng hơn.

3. Tại sao phải tinh chỉnh Flamingo?

Để kiểm chứng hiệu năng của hệ thống H100 mới, chúng tôi tiến hành thử nghiệm khả năng vận hành một mô hình ngôn ngữ lớn (LLMs). Trong quá trình đánh giá này, chúng tôi chọn tinh chỉnh một phiên bản Flamingo do cộng đồng phát triển. Dự án này đồng thời phục vụ hai mục tiêu: 

• System Validation: Chúng tôi sử dụng tác vụ tinh chỉnh này để kiểm thử nghiêm ngặt hạ tầng H100, nhằm đảm bảo nó có thể đáp ứng các yêu cầu tính toán khắt khe trong việc huấn luyện và vận hành một mô hình lớn. 

• Code Verification: Do mã nguồn gốc của mô hình Flamingo không được công bố rộng rãi, chúng tôi dựa vào phiên bản do cộng đồng phát triển. Quá trình này giúp kiểm chứng xem bản mã nguồn mở đó có tái hiện chính xác và vận hành được mô hình như mô tả trong bài báo nghiên cứu hay không. 

Do đó, xin lưu ý rằng trọng tâm của chúng tôi ở đây là khả năng của hệ thống, chứ không phải đánh giá độ chính xác của mô hình. 

Trong dự án này, một bản sao Flamingo có tên OpenFlamingo, do ML-Foundation phát triển, đã được sử dụng vì mô hình Flamingo gốc chưa được công bố rộng rãi. Mục tiêu là tinh chỉnh OpenFlamingo trên bộ dữ liệu gốc và đánh giá hiệu năng của nó trong điều kiện kiểm soát.  

Những kết quả này giúp xác định tính khả thi của việc triển khai OpenFlamingo trong các ứng dụng thực tiễn, đồng thời tối ưu hóa việc khai thác phần cứng.

4. Flamingo được tinh chỉnh như thế nào?

Cài đặt OpenFlamingo 

Để cài đặt gói này trong môi trường hiện có, hãy chạy lệnh sau: 

pip install open-flamingo

Hoặc để tạo một môi trường Conda dành cho việc chạy OpenFlamingo, hãy chạy lệnh sau:

conda env create -f environment.yml

Để cài đặt các dependencies phục vụ huấn luyện hoặc đánh giá, hãy chạy một trong hai lệnh đầu tiên. Để cài đặt toàn bộ, hãy chạy lệnh thứ ba.

pip install open-flamingo[training]
pip install open-flamingo[eval]
pip install open-flamingo[all]

Có ba tệp câu lệnh requirements.txt: 

• requirements.txt 

• requirements-training.txt 

• requirements-eval.txt  

Tùy theo mục đích sử dụng, bạn có thể cài đặt bất kỳ gói nào trong số này bằng lệnh: pip install -r <requirements-file.txt>. Tệp base chỉ bao gồm các dependencies cần thiết để chạy mô hình. 

Phát triển  

Các tác giả mã nguồn mở sử dụng pre-commit hooks để đồng bộ định dạng mã với các kiểm tra trong kho lưu trữ.

Để cài đặt pre-commit, hãy chạy lệnh sau: 

pip install pre-commit

hoặc sử dụng brew cho MacOS: 

 
brew install pre-commit  

Kiểm tra phiên bản đã cài đặt với: 

 
pre-commit --version  

Sau đó, tại thư mục gốc của kho lưu trữ này, hãy chạy lệnh: 

pre-commit install 

Sau đó, mỗi lần chúng ta chạy git commit, các kiểm tra sẽ được thực hiện. Nếu các tệp bị định dạng lại bởi các hooks, hãy chạy lại lệnh: 

git add    

cho các tệp bạn đã thay đổi, và hãy chạy lại lệnh 

git commit 

Quy trình huấn luyện  

Để huấn luyện OpenFlamingo, hãy đảm bảo môi trường của bạn khớp với môi trường của environment.yml. 

Xử lý dữ liệu  

Mã nguồn sử dụng WebDataset để tải các tệp .tar chứa các chuỗi hình ảnh và văn bản một cách hiệu quả. Chúng tôi khuyến nghị lấy mẫu lại các shard với thay thế trong quá trình huấn luyện bằng cách gắn cờ -- dataset_resampled. 

• Bộ dữ liệu LAION-2B 

LAION-2B chứa 2 tỷ cặp (hình ảnh, văn bản) được trích xuất từ web. Vui lòng sử dụng img2dataset để tải xuống bộ dữ liệu này thành các tệp tar. 

• Bộ dữ liệu đa phương thức C4 (Multimodal C4 - MMC4) 

OpenFlamingo được huấn luyện trên phiên bản đầy đủ của MMC4, bao gồm 103 triệu tài liệu với chuỗi hình ảnh và văn bản xen kẽ. Trong quá trình huấn luyện, chuỗi được cắt ngắn còn 256 token văn bản và 6 hình ảnh mỗi chuỗi. Codebase yêu cầu các tệp .tar chứa các tệp .json, trong đó bao gồm hình ảnh gốc được mã hóa bằng base64.  

Các script được cung cấp để chuyển đổi MMC4 sang định dạng này:  

(1) Tải xuống các shard MMC4 thành các tệp .zip bằng các script do MMC4 cung cấp (ví dụ: fewer_facesv2.sh).  

(2) Tải xuống hình ảnh gốc MMC4 vào một thư mục hình ảnh bằng các script do MMC4 cung cấp (ví dụ: download_images.py).  

(3) Chạy scripts/convert_mmc4_to_wds.py để chuyển đổi các mục đã tải xuống thành các tệp định dạng  tar. 

• Bộ dữ liệu tùy chỉnh 

Gần đây có báo cáo rằng các URL tải xuống bộ dữ liệu MMC4 đang gặp một số vấn đề về quyền truy cập. Do đó, chúng tôi đã tạo một script giúp chuẩn bị bộ dữ liệu tùy chỉnh bằng cách chuyển đổi nó sang định dạng của MMC4 (chúng tôi đã sử dụng bộ dữ liệu ADNI làm mục tiêu cho ví dụ này, với dữ liệu hình ảnh base64 mẫu cố định). Bạn có thể sửa đổi script này theo bộ dữ liệu tùy chỉnh của mình: 

import json
import os
import tarfile

def compress_directory_to_tar(directory_path):
    json_files = [f for f in os.listdir(directory_path) if f.endswith('.json')]
    os.makedirs('replicate_mmc4', exist_ok=True)

    for i in range(0, len(json_files), 20):
        batch_files = json_files[i:i+20]
        tar_file_path = os.path.join('replicate_mmc4', f"{i//20:09d}.tar")

        with tarfile.open(tar_file_path, "w:gz") as tar:
            for file in batch_files:
                tar.add(os.path.join(directory_path, file), arcname=file)

        print(f"Batch {i//20} compressed to {tar_file_path}")

def convert_adni_to_mmc4(input_json_path, output_folder):
    # Ensure the output folder exists
    os.makedirs(output_folder, exist_ok=True)

    # Load the large JSON file
    with open(input_json_path, 'r') as f:
        data = json.load(f)

    matched_text_index = 0

    # Iterate over each item in the list and save it as a separate JSON file
    for idx, item in enumerate(data):
        # Ensure compatibility with the structure of f9773b9c866145c28fe0b701dde8dfbe.json

        # Handle text list:
        conversations = item.get("conversations", None)
        if conversations is not None:
            text_list = []
            for conversation in conversations:
                text_list.append(conversation["value"])

        # Check for <image> tag in the first element of conversations list
        first_convo = conversations[0]["value"]
        if "<image>" in first_convo:
            if first_convo.startswith("<image>"):
                matched_text_index = 0
            elif first_convo.endswith("<image>"):
                matched_text_index = 1

        item["text_list"] = text_list

        # Handle image's base64 content:
        with open('./sample_base64.txt', 'r') as f:
            sample_img_base64_data = f.read()

        # Handle image info:
        img_info = []
        images_list = item.get("image", None)
        if images_list is not None:
            for img in images_list:
                img_obj = {}
                img_obj["image_name"] = img
                img_obj["raw_url"] = "https://example.com/{}".format(img)
                img_obj["matched_text_index"] = matched_text_index
                img_obj["matched_sim"] = 0.75
                img_obj["image_base64"] = sample_img_base64_data
                img_info.append(img_obj)

        # Create similarity_matrix
        similarity_matrix = []
        for img in img_info:
            for _ in range(len(text_list)):
                inner_list = [0] * len(text_list)
                inner_list[matched_text_index] = 1
                similarity_matrix.append(inner_list)

        # item["similarity_matrix"] = similarity_matrix

        output_item = {
            "id": item.get("id", None),
            "url": "https://example.com",
            "text_list": item.get("text_list", None),
            "image_info": img_info,
            "similarity_matrix": similarity_matrix,
            "could_have_url_duplicate": 0
        }

        # Save the item as a separate JSON file
        output_path = os.path.join(output_folder, f"{idx:05d}.json")
        with open(output_path, 'w') as out_f:
            json.dump(output_item, out_f)

Các chuỗi do ChatGPT khởi tạo 

Một tập con các mô hình của chúng tôi (liệt kê bên dưới) cũng được huấn luyện trên các chuỗi (hình ảnh, văn bản) do ChatGPT tạo thử nghiệm, trong đó hình ảnh được lấy từ LAION. Các shard chứa những chuỗi này có thể tìm thấy tại bảng tính CodaLab này. Họ không thể phân phối hình ảnh gốc trong các shard phát hành; hình ảnh cần được tải trước từ các URL trong các tệp JSON và chuyển sang định dạng base64 trước khi sử dụng dữ liệu này để huấn luyện trong codebase của chúng tôi. 

Các mô hình được huấn luyện với các chuỗi do ChatGPT tạo: 

• OpenFlamingo-4B-vitl-rpj3b 

• OpenFlamingo-4B-vitl-rpj3b-langinstruct 

Lệnh huấn luyện  

Một Slurm mẫu được cung cấp trong script huấn luyện trong scripts/. Bạn cũng có thể sửa đổi lệnh sau: 

 
torchrun --nnodes=1 --nproc_per_node=8 open_flamingo/train/train.py \
    --lm_path anas-awadalla/mpt-1b-redpajama-200b \
    --tokenizer_path anas-awadalla/mpt-1b-redpajama-200b \
    --cross_attn_every_n_layers 1 \
    --dataset_resampled \
    --batch_size_mmc4 2 \
    --train_num_samples_mmc4 1000 \
    --workers=4 \
    --run_name OpenFlamingo-3B-vitl-mpt1b \
    --num_epochs 20 \
    --warmup_steps 1875 \
    --mmc4_textsim_threshold 0.24 \
    --mmc4_shards "modifications/VLM_ADNI_DATA/replicate_mmc4/{000000000..000000040}.tar" \
    --report_to_wandb

Mã nguồn của mô hình MPT-1B (base và instruct) không nhận tham số labels và cũng không tính trực tiếp hàm mất mát cross-entropy trong forward(). Chúng tôi khuyến nghị sử dụng một phiên bản các mô hình MPT-1B đã được chỉnh sửa, có thể tìm thấy tại đây. 

Huấn luyện phân tán  

Theo mặc định, train.py sử dụng Distributed Data Parallel của Pytorch để huấn luyện.  

Để sử dụng Fully Sharded Data Parallel (FSDP), hãy gắn cờ --fsdp.  

Một số lưu ý về FSDP từ nhóm OpenFlamingo: Chúng tôi khuyến nghị sử dụng cờ --fsdp_use_orig_params df. Nếu --fsdp được bật mà không có cờ này, tất cả embeddings của mô hình ngôn ngữ sẽ được mở khóa trong quá trình huấn luyện. (Ngược lại, hành vi mặc định là chỉ huấn luyện các token <image> và <|endofchunk|> mới được thêm vào.)  

Lưu ý: Chúng tôi đã gặp vấn đề khi sử dụng OPT với cờ này. Các mô hình ngôn ngữ khác sẽ tương thích hơn.  

Chiến lược đóng gói FSDP hiện tại không cho phép huấn luyện embeddings của mô hình ngôn ngữ sử dụng tied weights (tức là embeddings đầu vào/đầu ra được liên kết). Để huấn luyện các mô hình như vậy với FSDP, embeddings của mô hình ngôn ngữ phải được đóng băng bằng cờ —freeze_lm_embeddings. 

 Chúng tôi cũng triển khai gradient checkpointing và huấn luyện với mixed precision. Sử dụng lần lượt các cờ —gradient_checkpointing và —precision. 

Khởi tạo mô hình OpenFlamingo  

OpenFlamingo hỗ trợ các bộ mã hóa thị giác đã được huấn luyện trước từ gói OpenCLIP, bao gồm các mô hình đã được huấn luyện trước của OpenAI. Chúng cũng hỗ trợ các mô hình ngôn ngữ đã được huấn luyện sẵn từ gói transformers, chẳng hạn như MPT, RedPajama, LLaMA, OPT, GPT-Neo, GPT-J hay các mô hình Pythia

 
from open_flamingo import create_model_and_transforms

model, image_processor, tokenizer = create_model_and_transforms(
    clip_vision_encoder_path="ViT-L-14",
    clip_vision_encoder_pretrained="openai",
    lang_encoder_path="anas-awadalla/mpt-1b-redpajama-200b",
    tokenizer_path="anas-awadalla/mpt-1b-redpajama-200b",
    cross_attn_every_n_layers=1,
    cache_dir="PATH/TO/CACHE/DIR" # Defaults to ~/.cache
) 

5. Kết quả đạt được

Dưới đây là kết quả được báo cáo từ WandBs của chúng tôi: GPU NVIDIA H100 

Hệ thống NVIDIA H100 đã được sử dụng:  

• Hệ thống được trang bị 8 * GPU NVIDIA H100 80GB HBM3. Tuy nhiên, đối với thiết lập huấn luyện này, chỉ cần 2 GPU với huấn luyện phân tán là hoàn thiện.  

• Mỗi GPU NVIDIA H100 có 80GB bộ nhớ băng thông cao (HBM3), tạo ra một hệ thống tính toán hiệu năng cao (HPC) hoặc huấn luyện AI.  

• Các GPU NVIDIA H100 ở trạng thái hiệu suất P0, cho thấy chúng đang ở chế độ hiệu suất cao nhất hiện có. 

Các số liệu được báo cáo của mô hình

Dữ liệu huấn luyện cho thấy quá trình đang vận hành tốt, với các hành vi đúng như yêu cầu trên các thông số khác nhau. Đường cong loss giảm mạnh ở giai đoạn đầu rồi dần ổn định, cho thấy mô hình đang đạt kết quả tốt. Learning rate tuân theo lịch trình khởi động tuyến tính (linear warm-up). Đây chính là phương pháp phổ biến để giữ cho giai đoạn huấn luyện ban đầu ổn định. 

Thời gian xử lý từng bước (step time) và thời gian tải dữ liệu nhìn chung ổn định, chỉ có một vài đột biến có thể xuất phát từ biến động hệ thống, quá trình lưu checkpoint hoặc độ trễ khi tải dữ liệu. Global step tăng tuyến tính, xác nhận rằng quá trình huấn luyện đang tiến triển ổn định. 

Chỉ số samples per second per GPU cũng duy trì mức ổn định, chỉ có một chút sụt giảm nhỏ nhưng không gây ảnh hưởng đáng kể đến hiệu năng tổng thể. Nhìn chung, các chỉ số này cho thấy mô hình huấn luyện hoạt động bình thường, tuy nhiên việc giám sát kỹ hơn những đột biến về thời gian xử lý bước và tải dữ liệu có thể giúp tối ưu hóa hiệu quả hơn nữa. 

Các số liệu được báo cáo của hệ thống 

1. Lỗi bộ nhớ không được sửa của GPU (Trên cùng bên trái): Đường này vẫn ở mức 0, cho thấy không có lỗi bộ nhớ nào chưa được sửa. 
2. Lỗi bộ nhớ đã được sửa của GPU (Trên cùng ở giữa): Biểu đồ cũng thẳng ở mức 0, nghĩa là không có lỗi bộ nhớ đã được sửa. 
3. Tốc độ xung nhịp bộ nhớ GPU (Trên cùng bên phải): Bình thường; tốc độ xung nhịp ổn đỉnh cho thấy không có sự thay đổi tần số động hay giới hạn hiệu năng. 
4. Tốc độ xung nhịp Bộ đa xử lý Streaming (SM) của GPU (Dưới cùng bên trái): Bình thường; tốc độ xung nhịp ổn định cho thấy không có điều tiết nhiệt. 
5. Mức sử dụng điện năng GPU (W) (Dưới cùng ở giữa): Hiển thị dạng chu kỳ, cho thấy mức tiêu thụ điện năng của GPU dao động trong quá trình thực thi tác vụ ⇒ có thể do xử lý theo nhóm, lập lịch công việc hoặc cơ chế quản lý điện năng.
6. Giới hạn điện năng cưỡng bức của GPU (W) (Trên cùng bên trái): Bình thường; điều này cho thấy GPU không vượt quá giới hạn điện năng được xác định trước.
7. Bộ nhớ GPU được phân bổ (Bytes) (Trên cùng ở giữa): Việc phân bổ bộ nhớ vẫn ổn định nhưng đột ngột giảm vào cuối => Sự sụt giảm này là khi quá trình huấn luyện kết thúc. 
8. Bộ nhớ GPU được phân bổ (%) (Trên cùng bên phải): Bình thường, tương tự như Bộ nhớ GPU được phân bổ (Bytes). 
9. Thời gian GPU dành để truy cập bộ nhớ (%) (Dưới cùng bên trái): Tương quan với Mức sử dụng điện năng GPU (W) ở trên. 
10. Nhiệt độ GPU (°C) (Dưới cùng ở giữa): Tương quan với Mức sử dụng điện năng GPU (W) ở trên. 
11. Mức sử dụng GPU (%) (Dưới cùng bên phải): Tương quan với Mức sử dụng điện năng GPU (W) ở trên.