在 Amazon EC2 上微调 NVIDIA Nemotron Speech ASR 以进行领域自适应

FROM nvcr.io/nvidia/pytorch:24.01-py3
ENV DEBIAN_FRONTEND=noninteractive
ENV TZ=UTC

# Set the working directory
WORKDIR /app

# Install system dependencies
RUN apt-get update && apt-get install -y \
    libsox-fmt-mp3 \
    gnupg \
    && rm -rf /var/lib/apt/lists/*

# Install Cython (needed for NeMo)
RUN pip install Cython

# Clone the specified branch of the pytorch-lightning repository and install it
RUN git clone -b bug_fix https://github.com/athitten/pytorch-lightning.git && \
    cd pytorch-lightning && \
    PACKAGE_NAME=pytorch pip install -e .

# Install TransformerEngine for optimization
RUN git clone https://github.com/NVIDIA/TransformerEngine.git && \
    cd TransformerEngine && \
    git fetch origin 8c9abbb80dba196f086b8b602a7cf1bce0040a6a && \
    git checkout FETCH_HEAD && \
    git submodule init && git submodule update && \

{
  "audio_filepath": "/path/to/audio.wav",
  "duration": 5.2,
  "text": "The transcription with punctuation and capitalization."
}

model:
  sample_rate: 16000
  normalize_text: true
  symbols_to_keep: ["'"]
  encoder:
    _target_: nemo.collections.asr.modules.ConformerEncoder
    feat_in: "${model.preprocessor.features}"
    n_layers: 17
    d_model: 512
    subsampling:
      dw_striding
    subsampling_factor: 8
    subsampling_conv_channels: 256

model:
  model_defaults:
    tdt_durations: [0, 1, 2, 3, 4]
    num_tdt_durations: 5
  loss:
    loss_name: "tdt"
    tdt_kwargs:
      durations: "${model.max_duration}"
      sigma: 0.02
      omega: 0.1

preprocessor:
  _target_: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor
  sample_rate: "${model.sample_rate}"
  normalize: "per_feature"
  window_size: 0.025
  window_stride: 0.01
  features: 128
  n_fft: 512
  dither: 0.00001

train_ds:
  augmentor:
    speed:
      prob: 0.4
      min_speed_rate: 0.9
      max_speed_rate: 1.1
    spec_augment:
      _target_: nemo.collections.asr.modules.SpectrogramAugmentation
      freq_masks: 2
      time_masks: 10
      freq_width: 27
      time_width: 0.05

trainer:
  devices: 8
  num_nodes: 8
  strategy:
    _target_: "lightning.pytorch.strategies.DeepSpeedStrategy"
  stage: 2
  offload_optimizer: true
  partition_activations: true
  gradient_as_bucket_view: true
  cpu_checkpointing: true
  contiguous_gradients: true
  overlap_comm: true

def get_base_model(self, trainer):
    """Get the base model to start training from based on config settings."""
    pretrained_name = self.config.init_from_pretrained_model

    # Handle multi-GPU download efficiently
    num_ranks = trainer.num_devices * trainer.num_nodes
    if num_ranks > 1 and is_global_rank_zero():
        logging.info(f"Downloading pretrained model '{pretrained_name}' on main process")
        asr_model = ASRModel.from_pretrained(model_name=pretrained_name)
    else:
        # Wait for model download to complete on main process
        wait_time = 1 if is_global_rank_zero() else 60
        logging.info(f"Waiting {wait_time}s for model download")
        time.sleep(wait_time)
        asr_model = ASRModel.from_pretrained(model_name=pretrained_name)

    asr_model.to(f"cuda:{int(os.environ.get('LOCAL_RANK', 0))}")

    # Unfreezing encoders to update the parameters
    asr_model.encoder.unfreeze()
    logging.info("Model encoder has been un-frozen")

    return asr_model

def train(self, model_path):
    """Train the ASR model."""
    # Create trainer with experiment manager for logging
    trainer = self.create_trainer()

    # Initialize model based on config settings
    asr_model = self.get_base_model(trainer)

    # Setup dataloaders
    asr_model = self.setup_dataloaders(asr_model)

    # Setup optimization
    asr_model.setup_optimization(self.config.model.optim)

    # Setup SpecAug if available
    if hasattr(self.config.model, 'spec_augment'):
        asr_model.spec_augment = ASRModel.from_config_dict(self.config.model.spec_augment)

    # Train the model
    trainer.fit(asr_model)

    # Save the trained model
    asr_model.save_to(model_path)

    return model_path

exp_manager:
  create_mlflow_logger: true
  mlflow_logger_kwargs: {"tracking_uri": "mlruns"}
  checkpoint_callback_params:
    monitor: "val_wer"
    mode: "min"
    save_top_k: 10

trainer:
  devices: 8  # Number of GPUs to use (-1 would use all available)
  num_nodes: 8  # Number of compute nodes (servers) for distributed training
  max_epochs: 2  # Maximum number of training epochs
  max_steps: -1  # Maximum number of training steps (-1 means compute from epochs)
  val_check_interval: 1.0  # Validation frequency: 1.0 = once per epoch, 0.25 = 4 times per epoch
  accelerator: auto  # Hardware accelerator: auto, gpu, cpu
  strategy: "lightning.pytorch.strategies.DeepSpeedStrategy"  # "lightning.pytorch.strategies.DeepSpeedStrategy" or "lightning.pytorch.strategies.DDPStrategy"
  # Distributed training strategy
  stage: 2
  offload_optimizer: true  # Offload optimizer states to CPU to save GPU memory
  offload_parameters: false  # Don't offload model parameters to CPU
  partition_activations: true  # Partition activations across GPUs to save memory
  gradient_as_bucket_view: true  # Use bucket view for gradients to save memory
  cpu_checkpointing: true  # Store activations on CPU during backward pass
  contiguous_gradients: true  # Ensure contiguous memory for gradients
  overlap_comm: true  # Overlap communication and computation
  allgather_bucket_size: 2e8  # Bucket size for all-gather operations
  reduce_bucket_size: 2e8  # Bucket size for reduce operations
  zero_force_ds_cpu_optimizer: false  # Don't force CPU optimizer with ZeRO

from fastapi import FastAPI, UploadFile, File, BackgroundTasks
from typing import Optional, List

app = FastAPI(
    title="Optimized NeMo ASR OpenAI-Compatible API"
)

@app.post("/v1/audio/transcriptions")
async def create_transcription(
    background_tasks: BackgroundTasks,
    file: UploadFile = File(..., description="Audio file to transcribe"),
    model: Optional[str] = "nemo-parakeet-rnnt-1.1b",
    language: Optional[str] = "en",
    prompt: Optional[str] = None,
    response_format: Optional[str] = "json",
    temperature: Optional[float] = 0.0,
    timestamp_granularities: Optional[List[str]] = None
):

FROM amazonlinux:2023

# Install Python and essential tools
RUN dnf update -y && \
    dnf install -y shadow-utils python3.11 python3.11-pip python3.11-devel gcc gcc-c++ wget tar gzip xz && \
    dnf clean all && \
    ln -sf /usr/bin/python3.11 /usr/bin/python && \
    ln -sf /usr/bin/pip3.11 /usr/bin/pip

# Install ffmpeg from static build (recommended for AL2023)
RUN wget -q https://johnvansickle.com/ffmpeg/releases/ffmpeg-release-amd64-static.tar.xz && \
    tar -xf ffmpeg-release-amd64-static.tar.xz && \
    mv ffmpeg-*-amd64-static/ffmpeg /usr/local/bin/ && \
    mv ffmpeg-*-amd64-static/ffprobe /usr/local/bin/ && \
    rm -rf ffmpeg-* && \
    chmod +x /usr/local/bin/ffmpeg /usr/local/bin/ffprobe

# Some parts od Dockerfile is skipped in here. Refer to repo for a full file
RUN pip install --no-cache-dir --upgrade pip && \
    pip install --no-cache-dir fastapi uvicorn python-multipart soundfile uvloop prometheus-fastapi-instrumentator

RUN pip install -v --no-cache-dir nemo_toolkit[asr] torch cuda-python>=12.3