Загрузка данных в среде выполнения ИИ

Это важно

Эта функция доступна в общедоступной предварительной версии.

Ресурсы данных и моделей имеют решающее значение для глубокого обучения и задач постобучения больших языковых моделей (LLM) и визуально-языковых моделей (VLM). В AI Runtime все данные и ресурсы моделей доступны через Unity Catalog:

Тома Unity Catalog: используются преимущественно для больших наборов данных и неструктурированных файлов, включая изображения, аудио и текст.
Таблицы Unity Catalog: используются для структурированных и табличных данных, доступны через Spark Connect.

Ваши тома и таблицы должны быть зарегистрированы в каталоге Unity и доступны вашему пользователю или руководителю сервиса.

Том Unity Catalog для неструктурированных данных

Тома каталога Unity обеспечивают регулируемый доступ к нетабличным данным в любом формате, включая структурированные, полуструктурированные и неструктурированные данные. В AI Runtime тома являются основным механизмом доступа к большим наборам данных, тексту, активам моделей и контрольным точкам моделей.

Пользователи могут перечислять, читать и записывать файлы в томах Unity Catalog, используя знакомые операции с файловой системой, аналогично работе с файлами на локальном диске:

import os

dir_path = "/Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir"
file_path = os.path.join(dir_path, "test_file")

os.makedirs(dir_path, exist_ok=True)

# Write to the file
with open(file_path, "w") as file:
    file.write("Hello, World!")

Аналогично, операции с оболочками работают аналогично:

%sh ls -l /Volumes/<catalog-name>/<schema-name>/<volume-name>
%sh mkdir -p /Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir
%sh touch /Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir/test_file

Некоторые характеристики томов Unity Catalog делают их отлично подходящими для задач машинного обучения:

Распределённое хранилище: Unity Catalog поддерживается распределённым хранилищем, позволяя AI Runtime читать и записывать данные и моделировать активы по всей платформе, как из ноутбуков, так и с CLI-ориентированных рабочих нагрузок.
Оптимизировано для схем доступа в ML: Базовое хранилище и пути доступа оптимизированы для типичных ML-нагрузок, особенно для больших файлов с последовательным чтением и записью. Это делает Unity Catalog хорошо подходящим для загрузки обучающих данных, загрузки артефактов модели и записи контрольных точек модели.
Доступ, подобный файловой системе: пользователи могут перечислять, читать и записывать файлы в томах Unity Catalog, используя привычные операции файловой системы, аналогично работе с файлами на локальном диске.

Благодаря автоматическим фоновым коммитам пользователи могут рассчитывать на стабильный доступ к данным томов Unity Catalog:

Записывает: AI Runtime автоматически фиксирует записи, делая изменения видимыми для других приложений и рабочих нагрузок, обращающихся к тому же тому же тому Unity Catalog.
Чтение: AI Runtime автоматически обнаруживает изменения в томе без необходимости выполнять явное обновление или синхронизацию со стороны пользователя.

Производительность настройки громкости

Как уже упоминалось, тома Unity Catalog поддерживаются распределённым хранилищем и оптимизированы для больших файлов с последовательным чтением и записью.

Несколько советов, которые помогут вам получить наилучшую производительность от AI Runtime:

Объединяйте данные в более крупные файлы: По возможности объединяйте их в меньшее, более крупное количество файлов — примерно от 1 ГиБ до 10 ГиБ на файл. Это позволяет AI Runtime агрессивно предзагружать данные и автоматически достигать почти оптимальной последовательной производительности чтения.
Для нагрузок с малыми файлами используйте локальный диск: если ваша нагрузка связана с большим количеством мелких файлов, рассмотрите возможность скопирования файлов на локальный диск с помощью параллельных копий перед обработкой. Это может сократить накладные расходы при многократном доступе к множеству небольших файлов через том.
```
# Recommended using parallel copy (256 concurrency in this example, you can tune)
#
# This takes only 22 seconds to copy 15,375 150KiB small image files.
%sh cd /Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir/ && find . -type f -print0 | xargs -0 -P 256 -I {} cp --parents "{}" /tmp/


# !!! Avoid doing this !!!
#
# Because the files are copied in serial, this copies the same 15,375 150KiB small image files much more slowly.
# %sh cp -r /Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir/* /tmp
```
Вы можете использовать UCVolumeDataset это для своих задач по машинному обучению. Он включает описанные выше оптимизации для эффективного доступа к данным и загрузки из томов Unity Catalog. Дополнительные сведения см. в следующих разделах.

Загрузите неструктурированные данные с помощью `UCVolumeDataset`

Для неструктурированных данных, таких как изображения, аудио и текстовые файлы, хранящиеся в томах каталога Unity, используйте UCVolumeDataset из serverless_gpu.data пакета. UCVolumeDataset — это PyTorch IterableDataset , копирующий каждый файл из тома в быстрый локальный кэш при первом доступе и возвращающий кэшированный локальный путь к файлу. Он обрабатывает проблемы производительности и распределения, которые вы бы в противном случае реализовали вручную:

Локальное кэширование. При первом обращении файлы копируются из точки монтирования FUSE в локальный каталог кэша, а затем считываются из кэша, поэтому обучение в течение нескольких эпох не требует повторного чтения тома.
Автоматическое секционирование. При инициализации torch.distributed файлы разделяются между рангами, а затем дополнительно между рабочими процессами DataLoader, поэтому каждая пара (rank, worker) получает непересекающийся фрагмент без дополнительной конфигурации.

Замечание

UCVolumeDataset и serverless_gpu.data.DataLoaderтребуется среда GPU 5 или более поздней версии.

UCVolumeDataset возвращает необработанные пути к локальным файлам. Чтобы декодировать эти файлы в тензоры, оберните их вторым IterableDataset, который использует поток путей и применяет вашу логику разбора. Это позволяет разделить аспекты ввода-вывода и синтаксического разбора.

from serverless_gpu.data import UCVolumeDataset
from torch.utils.data import IterableDataset
from PIL import Image
import torchvision.transforms.functional as TF

class ImageDataset(IterableDataset):
    """Decodes each cached file path from UCVolumeDataset into a tensor."""

    def __init__(self, path_dataset: UCVolumeDataset):
        self._path_dataset = path_dataset

    def __iter__(self):
        for local_path in self._path_dataset:
            image = Image.open(local_path).convert("RGB")
            yield TF.to_tensor(image)

path_dataset = UCVolumeDataset("/Volumes/catalog/schema/my_volume/images")
dataset = ImageDataset(path_dataset)

Обёртка получает уже кэшированные локальные пути, поэтому этап разбора никогда не касается тома. Можно привязать дополнительные оболочки для расширения, маркеризации или фильтрации.

Для оптимальной производительности используйте UCVolumeDataset вместе с serverless_gpu.data.DataLoader, а не стандартный DataLoader из PyTorch. Он настроен на AI Runtime I/O и одновременно загружает и кэширует файлы, пока GPU вычисляет данные.

Модели контрольных точек на объёмах

Чтобы проверить вашу модель и возобновить обучение с последнего снимка или восстановиться после сбоя, вы можете использовать тома Unity Catalog так же, как локальную файловую систему.

Databricks рекомендует использовать распределённую контрольную точку (DCP) для лучшей производительности как на одной GPU, так и на нескольких GPU. См. «Быстрое, отказоустойчивое обучение PyTorch по AI Runtime » из инженерного блога Databricks.

import torch.distributed.checkpoint as dcp
from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict

import serverless_gpu

checkpoint_path = "/Volumes/my-catalog/my-schema/my-volume/checkpoints/step_1000"

# Save
model_sd, optim_sd = get_state_dict(model, optimizer)
state_dict = {"model": model_sd, "optim": optim_sd, "step": 1000}
dcp.async_save(
    state_dict,
    storage_writer=serverless_gpu.data.UCVolumeWriter(checkpoint_path))

# Load
model_sd, optim_sd = get_state_dict(model, optimizer)
state_dict = {"model": model_sd, "optim": optim_sd}
dcp.load(
    state_dict,
    storage_reader=serverless_gpu.data.UCVolumeReader(checkpoint_path))

set_state_dict(
    model,
    optimizer,
    model_state_dict=state_dict["model"],
    optim_state_dict=state_dict["optim"],
)

Монолитный подход с использованием torch.save тоже работает.

Для сохранения контрольных точек модели с одним GPU,

# The monolithic torch.save approach for single GPU chip

# Save
torch.save({"model": model.state_dict(), "opt": optimizer.state_dict()},
           "/Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir/ckpt.pt")


# Load
ckpt = torch.load(
    "/Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir/ckpt.pt",
    weights_only=True)
model.load_state_dict(ckpt["model"])
optimizer.load_state_dict(ckpt["opt"])

Для распределённого обучения, запускаемого через torchrun,

# The monolithic torch.save approach for multi-GPU distributed training.
# This snippet assumes your launcher has already called
# dist.init_process_group(...).

import os
import torch.distributed as dist

# Save only on rank 0.
if dist.get_rank() == 0:
    torch.save({"model": model.state_dict(), "opt": optimizer.state_dict()},
               "/Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir/ckpt.pt")

# Wait for rank 0 to finish writing before any rank reads.
dist.barrier()

# Load on ALL ranks (map to current rank's local GPU).
local_rank = int(os.environ["LOCAL_RANK"])
ckpt = torch.load(
    "/Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir/ckpt.pt",
    map_location=f"cuda:{local_rank}",
    weights_only=True)
model.load_state_dict(ckpt["model"])
optimizer.load_state_dict(ckpt["opt"])

Загрузка табличных данных

Используйте Spark Connect для загрузки табличных данных машинного обучения из таблиц Delta.

Для обучения с одним узлом можно преобразовать Apache Spark DataFrames в pandas DataFrames с помощью метода PySpark, а затем при необходимости преобразовать в формат NumPy с помощью метода toPandas()to_numpy().

Замечание

Spark Connect откладывает анализ и разрешение имен во время выполнения, что может изменить поведение кода. См . статью "Сравнение Spark Connect с классической версией Spark".

Spark Connect поддерживает большинство API PySpark, включая Spark SQL, Pandas API в Spark, Structured Streaming и MLlib (на базе DataFrame). См. справочную документацию по API PySpark для последних поддерживаемых API.

Сведения о других ограничениях см. в разделе об ограничениях бессерверных вычислений.

Загрузите большие таблицы Delta с помощью томов Unity Catalog

Для больших таблиц Delta, которые слишком большие для преобразования toPandas(), экспортируйте данные в том каталога Unity и загружайте их непосредственно с помощью PyTorch или Hugging Face:

# Step 1: Export the Delta table to Parquet files in a UC volume
output_path = "/Volumes/catalog/schema/my_volume/training_data"
spark.table("catalog.schema.my_table").write.mode("overwrite").parquet(output_path)

# Step 2: Load the exported data directly using Hugging Face datasets
from datasets import load_dataset

dataset = load_dataset("parquet", data_files="/Volumes/catalog/schema/my_volume/training_data/*.parquet")

Этот подход позволяет избежать накладных расходов Spark во время обучения и хорошо подходит как для рабочих процессов обучения с одним GPU, так и для распределенных рабочих процессов.

Обратная связь

Были ли сведения на этой странице полезными?

Last updated on 2026-09-30