Распределенное обучение с TensorFlow и Python

strategy = tf.distribute.MirroredStrategy()
with strategy.scope():
    model = create_model()  # Стандартное создание модели
    model.compile(...)      # Стандартная компиляция
    
# И даже обучение выглядит как обычно!
model.fit(dataset, epochs=10)

# Пример TF_CONFIG для узла worker с индексом 1
os.environ["TF_CONFIG"] = json.dumps({
    "cluster": {
        "worker": ["10.0.0.1:8000", "10.0.0.2:8000"]
    },
    "task": {"type": "worker", "index": 1}
})

strategy = tf.distribute.MultiWorkerMirroredStrategy()
# И дальше как с MirroredStrategy

resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='grpc://10.0.0.1:8470')
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.TPUStrategy(resolver)

class CustomDistributionStrategy(tf.distribute.Strategy):
    def __init__(self, devices=None):
        super(CustomDistributionStrategy, self).__init__()
        # Собственная логика инициализации
        
    def _make_dataset_iterator(self, dataset):
        # Кастомная логика итерации по датасету
        pass
        
    def _make_input_fn_iterator(self, input_fn, replication_mode):
        # Логика для работы с input_fn
        pass

# Оптимизированный пайплайн данных для распределенного обучения
def make_dataset(batch_size):
    dataset = tf.data.Dataset.from_tensor_slices(...)
    # Кэшируем после предобработки
    dataset = dataset.map(preprocess_fn, num_parallel_calls=tf.data.AUTOTUNE).cache()
    # Перемешиваем с достаточно большим буфером
    dataset = dataset.shuffle(buffer_size=10000)
    # Важно: глобальный батч-сайз будет batch_size * число_реплик
    dataset = dataset.batch(batch_size)
    # Критически важно для производительности!
    dataset = dataset.prefetch(tf.data.AUTOTUNE)
    return dataset

# Динамическая балансировка с помощью tf.data и теневого режима
def create_balanced_dataset(files, num_workers):
    dataset = tf.data.Dataset.from_tensor_slices(files)
    # Теневой режим: устройства запрашивают следующий батч, 
    # когда готовы обработать предыдущий
    options = tf.data.Options()
    options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.DATA
    dataset = dataset.with_options(options)
    return dataset.batch(batch_size).prefetch(tf.data.AUTOTUNE)

class DeviceAwareBalancer:
    def __init__(self, devices_perf_coeffs):
        # Словарь {device_name: relative_performance}
        self.device_performance = devices_perf_coeffs
        total_perf = sum(devices_perf_coeffs.values())
        self.device_ratios = {d: p/total_perf for d, p in devices_perf_coeffs.items()}
    
    def get_batch_sizes(self, global_batch):
        # Распределяем примеры пропорционально производительности
        return {d: int(global_batch * r) for d, r in self.device_ratios.items()}

cross_device_ops = tf.distribute.HierarchicalCopyAllReduce(
    num_packs=2,  # Разбиваем градиенты на пакеты
    collective_keys=collective_keys
)
strategy = tf.distribute.MultiWorkerMirroredStrategy(
    communication_options=tf.distribute.experimental.CommunicationOptions(
        implementation=tf.distribute.experimental.CommunicationImplementation.NCCL,
        collective_ops_options=collective_ops_options
    )
)

# Синхронная стратегия - все ждут завершения вычислений на всех устройствах
strategy = tf.distribute.MirroredStrategy()
with strategy.scope():
    model = create_model()
    opt = tf.keras.optimizers.Adam(learning_rate=0.01)
    model.compile(optimizer=opt, loss='sparse_categorical_crossentropy')
    
    # Каждый шаг обучения - синхронная операция
    model.fit(dataset, epochs=5)

# Асинхронная стратегия - каждое устройство обновляет веса независимо
strategy = tf.distribute.experimental.ParameterServerStrategy(
    cluster_resolver=cluster_resolver  # Настройки кластера с параметр-серверами
)
with strategy.scope():
    model = create_model()
    # Важно! Асинхронный оптимизатор 
    opt = tf.keras.optimizers.experimental.SGD(learning_rate=0.01)
    model.compile(optimizer=opt, loss='sparse_categorical_crossentropy')
    
    model.fit(dataset, epochs=5)

# Псевдокод для группованной синхронизации
def train_step(inputs):
    features, labels = inputs
    
    # Записываем текущую версию модели
    current_version = get_model_version()
    
    with tf.GradientTape() as tape:
        predictions = model(features, training=True)
        loss = loss_fn(labels, predictions)
    
    gradients = tape.gradient(loss, model.trainable_variables)
    
    # Проверяем, не устарела ли наша версия модели слишком сильно
    if get_model_version() - current_version < MAX_STALENESS:
        # Применяем градиенты
        optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    
    return loss

class MyCustomStrategy(tf.distribute.Strategy):
  def __init__(self):
    super(MyCustomStrategy, self).__init__()
    # Инициализация специфичных параметров
    
  def _create_variable(self, next_creator, **kwargs):
    # Определяет, как создаются переменные
    return next_creator(**kwargs)
    
  def _distribute_dataset(self, dataset_fn):
    # Как распределяются данные между устройствами
    return dataset_fn()
    
  def _run(self, fn, args, kwargs):
    # Как запускается функция на каждом устройстве
    return self.extended.call_for_each_replica(fn, args=args, kwargs=kwargs)

class AdaptiveBatchStrategy(tf.distribute.Strategy):
  def __init__(self, device_performance_map):
    """
    Args:
      device_performance_map: словарь {имя_устройства: относительная_производительность}
    """
    self._device_performance = device_performance_map
    # Нормализуем значения производительности
    total = sum(device_performance_map.values())
    self._device_ratios = {k: v/total for k, v in device_performance_map.items()}
    
    # Список всех устройств
    self._devices = list(device_performance_map.keys())
    
    # Создаем контекст для каждого устройства
    self._device_contexts = {d: tf.device(d) for d in self._devices}
 
  def _distribute_dataset(self, dataset_fn):
    # Для каждого устройства создаем подходящий по размеру датасет
    per_device_datasets = {}
    main_dataset = dataset_fn()
    
    for device, ratio in self._device_ratios.items():
      # Батч подстраивается под мощность устройства
      adjusted_batch = max(1, int(global_batch_size * ratio))
      device_dataset = main_dataset.batch(adjusted_batch)
      per_device_datasets[device] = device_dataset
      
    return tf.distribute.DistributedDataset(per_device_datasets)
    
  def _run(self, fn, args, kwargs):
    # Запускаем функцию на каждом устройстве с соответствующим контекстом
    results = {}
    for device in self._devices:
      with self._device_contexts[device]:
        results[device] = fn(*args, **kwargs)
    
    # Агрегируем результаты с учетом мощности устройств
    # (для тренировочного цикла это обычно градиенты)
    aggregated_result = self._aggregate_results(results)
    return aggregated_result

class EnhancedMirroredStrategy(tf.distribute.MirroredStrategy):
  def __init__(self, devices=None, cross_device_ops=None):
    super(EnhancedMirroredStrategy, self).__init__(devices, cross_device_ops)
    
  # Переопределяем только нужные методы, остальная логика наследуется
  def _distribute_dataset(self, dataset_fn):
    dataset = super()._distribute_dataset(dataset_fn)
    # Добавляем дополнительную логику обработки датасета
    return dataset

# Типичная настройка TPU в TensorFlow
resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='grpc://10.0.0.1:8470')
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
 
# Создаем TPU-специфичную стратегию
strategy = tf.distribute.TPUStrategy(resolver)
 
# Далее работаем как обычно, но в рамках контекста стратегии
with strategy.scope():
    model = create_model()
    model.compile(optimizer='adam', loss='categorical_crossentropy')
    
# Обучение запустится автоматически на всех доступных ядрах TPU
model.fit(dataset, epochs=10)

def create_tpu_optimized_dataset(global_batch_size, strategy):
    # Важно делать батч кратным числу ядер TPU
    cores = strategy.num_replicas_in_sync
    adjusted_batch = ((global_batch_size + cores - 1) // cores) * cores
    
    dataset = tf.data.Dataset.from_tensor_slices(...)
    # Важно! Drop remainder критичен для оптимальной работы TPU
    dataset = dataset.batch(adjusted_batch // cores, drop_remainder=True)
    # TPU любят предварительную загрузку данных
    dataset = dataset.prefetch(tf.data.AUTOTUNE)
    return dataset

# Подключение профилировщика в коде
tf.profiler.experimental.start('logdir')
# Код обучения модели
for epoch in range(epochs):
    for batch in dataset:
        train_step(batch)
    # Делаем снимок профиля через каждые несколько эпох
    if epoch % profile_every == 0:
        tf.profiler.experimental.stop()
        # Запускаем снова для следующего интервала
        tf.profiler.experimental.start('logdir')
 
# Останавливаем профилировщик в конце
tf.profiler.experimental.stop()

# Создаем колбэк для профилирования
class ProfilerCallback(tf.keras.callbacks.Callback):
    def __init__(self, log_dir, profile_batches):
        super(ProfilerCallback, self).__init__()
        self.log_dir = log_dir
        self.profile_batches = profile_batches
        
    def on_train_begin(self, logs=None):
        # Запускаем профилировщик при начале обучения
        tf.profiler.experimental.start(self.log_dir)
        
    def on_train_batch_end(self, batch, logs=None):
        # Периодически останавливаем и запускаем профилировщик
        if batch in self.profile_batches:
            tf.profiler.experimental.stop()
            print(f"Профиль сохранен для батча {batch}")
            # Небольшая пауза, чтобы профиль точно успел записаться
            time.sleep(1)
            tf.profiler.experimental.start(self.log_dir)
            
    def on_train_end(self, logs=None):
        # Останавливаем профилировщик в конце обучения
        tf.profiler.experimental.stop()
 
# Использование колбэка с моделью Keras
model.fit(dataset, epochs=10, callbacks=[
    ProfilerCallback(log_dir='./logs/profile', profile_batches=[100, 200, 300])
])

# Оптимизация пайплайна данных на основе профилирования
def create_optimized_dataset(files, batch_size):
    # Увеличиваем число параллельных процессов для чтения
    dataset = tf.data.Dataset.from_tensor_slices(files)
    dataset = dataset.interleave(
        lambda file: tf.data.TFRecordDataset(file),
        num_parallel_calls=tf.data.AUTOTUNE,  # Важный параметр!
        cycle_length=16  # Определяется экспериментально
    )
    # Кэшируем после тяжелых преобразований
    dataset = dataset.map(parse_fn, num_parallel_calls=tf.data.AUTOTUNE).cache()
    dataset = dataset.shuffle(10000)
    dataset = dataset.batch(batch_size)
    # Буфер предзагрузки критичен для производительности
    dataset = dataset.prefetch(tf.data.AUTOTUNE)
    return dataset

# Пример разделения модели на части для конвейерного параллелизма
strategy = tf.distribute.experimental.TPUStrategy(...)
 
with strategy.scope():
    # Создаем модель по частям
    encoder_layers = [TransformerBlock() for _ in range(num_layers//2)]
    decoder_layers = [TransformerBlock() for _ in range(num_layers//2)]
    
    # Размещаем части на разных устройствах
    with tf.device('/GPU:0'):
        input_embeddings = tf.keras.layers.Embedding(...)
        encoder_output = input_embeddings(inputs)
        for layer in encoder_layers:
            encoder_output = layer(encoder_output)
    
    with tf.device('/GPU:1'):
        decoder_output = encoder_output  # Передача между устройствами
        for layer in decoder_layers:
            decoder_output = layer(decoder_output)
        outputs = tf.keras.layers.Dense(...)(decoder_output)

# Разделение состояния оптимизатора
optimizer = tf.keras.optimizers.Adam(...)
sharded_optimizer = tf.keras.mixed_precision.experimental.LossScaleOptimizer(
    optimizer, dynamic=True
)
 
with strategy.scope():
    model = create_transformer_model()
    model.compile(optimizer=sharded_optimizer, ...)

# Включение gradient checkpointing для экономии памяти
model = create_transformer_model()
model = tf.keras.models.Model(model.inputs, model.outputs)
 
# Активируем checkpointing для экономии памяти за счет повторных вычислений
tf.keras.backend.clear_session()
tf.config.optimizer.set_experimental_options({'auto_mixed_precision': True})
tf.config.optimizer.set_jit(True)  # XLA-компиляция для эффективности
tf.config.experimental.enable_tensor_float_32_execution(True)  # Для NVIDIA A100

# Включаем тренировку с смешанной точностью
policy = tf.keras.mixed_precision.Policy('mixed_float16')
tf.keras.mixed_precision.set_global_policy(policy)
 
# Теперь модель будет использовать float16 для вычислений
model = create_transformer_model()

def get_strategy(worker_count, peak_load=False):
    if peak_load and worker_count > 4:
        # В пиковую нагрузку используем асинхронное обучение
        return tf.distribute.experimental.ParameterServerStrategy(...)
    else:
        # В обычном режиме - синхронное для лучшей сходимости
        return tf.distribute.MultiWorkerMirroredStrategy(...)
        
# Периодически проверяем загрузку системы
is_peak = monitoring_service.get_system_load() > THRESHOLD
strategy = get_strategy(available_workers, is_peak)

# Без Horovod
optimizer = tf.keras.optimizers.Adam(0.001)
 
# С Horovod (добавляется буквально пара строк)
import horovod.tensorflow as hvd
hvd.init()
optimizer = tf.keras.optimizers.Adam(0.001 * hvd.size())
optimizer = hvd.DistributedOptimizer(optimizer)

from ray import tune
 
# Определяем пространство поиска
config = {
  "learning_rate": tune.loguniform(1e-4, 1e-1),
  "batch_size": tune.choice([32, 64, 128]),
  # Остальные параметры...
}
 
# Запускаем параллельный поиск на кластере
analysis = tune.run(
  train_fn,  # Функция обучения
  config=config,
  resources_per_trial={"cpu": 2, "gpu": 1},
  num_samples=50  # Число экспериментов
)

# Эластичный запуск с Horovod
import horovod.tensorflow as hvd
from horovod.tensorflow.elastic import run
 
def train_fn():
    # Инициализация Horovod
    hvd.init()
    
    # Получаем текущее количество рабочих узлов
    current_workers = hvd.size()
    
    # Масштабируем learning rate с учетом изменяющегося количества узлов
    lr = base_lr * current_workers
    
    # Создаем модель и оптимизатор
    model = create_model()
    optimizer = tf.keras.optimizers.Adam(lr)
    optimizer = hvd.DistributedOptimizer(optimizer)
    
    # Остальной код обучения...
    
# Запуск с возможностью эластичного масштабирования
min_workers = 2  # Минимальное количество узлов для продолжения работы
max_workers = 10  # Максимально поддерживаемое количество
run(train_fn, min_workers=min_workers, max_workers=max_workers)

import ray
from ray import tune
import tensorflow as tf
 
# Инициализируем Ray
ray.init()
 
# Определяем функцию обучения модели
@ray.remote(num_gpus=1)
def train_model(config):
    # Настройка TensorFlow
    strategy = tf.distribute.MirroredStrategy()
    with strategy.scope():
        model = create_model(config)
        # Обучение модели...
        
    return final_metrics
 
# Параллельный запуск множества конфигураций
configs = [{"lr": 0.001, "batch_size": 32}, 
           {"lr": 0.01, "batch_size": 64},
           # Другие конфигурации...
          ]
results = ray.get([train_model.remote(config) for config in configs])

import flexflow
from flexflow.keras.models import Model
from flexflow.keras.layers import Input, Dense
 
# Определяем модель используя API, совместимый с Keras
input_tensor = Input(shape=(784,))
x = Dense(512, activation='relu')(input_tensor)
output_tensor = Dense(10, activation='softmax')(x)
model = Model(input_tensor, output_tensor)
 
# Компилируем с FlexFlow
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')
 
# Автоматическая оптимизация стратегии распределения
config = flexflow.config()
config.search_budget = 10000  # Количество стратегий для перебора
model.fit(x_train, y_train, epochs=5)

# Пример интеграции DeepSpeed с TensorFlow (экспериментальная поддержка)
import deepspeed
import tensorflow as tf
 
# Шардирование оптимизатора с помощью ZeRO
optimizer = deepspeed.DeepSpeedZeroOptimizer(
    tf.keras.optimizers.Adam(lr=3e-5),
    stage=2,  # Уровень оптимизации памяти (1, 2 или 3)
    contiguous_gradients=True
)
 
with strategy.scope():
    model = create_giant_model()
    model.compile(optimizer=optimizer, loss='mse')
    model.fit(dataset, epochs=10)