# NeMo Automodel

[NeMo Automodel](https://github.com/NVIDIA-NeMo/Automodel)은 NVIDIA에서 제공하는 PyTorch DTensor 네이티브 학습 라이브러리로, 대규모로 diffusion 모델을 fine-tuning 및 pretraining하는 데 사용됩니다. Hugging Face 네이티브로, Hub의 Diffusers 형식 모델을 checkpoint 변환 없이 학습할 수 있습니다. 동일한 YAML recipe와 수정 가능한 학습 스크립트가 1개 GPU부터 수백 개의 노드까지 모든 규모에서 실행되며, [FSDP2](https://pytorch.org/docs/stable/fsdp.html) 분산 학습, 다중 해상도 버킷 기반 dataloading, 최대 GPU 활용을 위한 사전 인코딩된 latent space 학습을 지원합니다. 학습에는 [flow matching](https://huggingface.co/papers/2210.02747)을 사용하며, 완전한 오픈 소스(Apache 2.0)로 NVIDIA의 지원을 받고 있으며 적극적으로 유지 관리되고 있습니다.

NeMo Automodel은 Diffusers와 직접 통합됩니다. Diffusers 모델 클래스를 사용하여 Hugging Face Hub에서 사전 학습된 모델을 로드하고, [DiffusionPipeline](/docs/diffusers/main/ko/api/pipelines/overview#diffusers.DiffusionPipeline)으로 출력을 생성합니다.

일반적인 워크플로우는 NeMo Automodel을 설치(pip 또는 Docker)하고, 데이터를 `.meta` 파일로 인코딩하여 준비한 후, YAML 레시피를 구성하고, `torchrun`으로 학습을 시작하고, 생성된 체크포인트로 추론을 실행하는 것입니다.

## 지원되는 모델

| 모델 | Hugging Face ID | 작업 | 매개변수 | 사용 사례 |
|-------|----------------|------|------------|----------|
| Wan 2.1 T2V 1.3B | [Wan-AI/Wan2.1-T2V-1.3B-Diffusers](https://huggingface.co/Wan-AI/Wan2.1-T2V-1.3B-Diffusers) | 텍스트-투-비디오 | 1.3B | 제한된 하드웨어에서 비디오 생성 (단일 40GB A100에 적합) |
| FLUX.1-dev | [black-forest-labs/FLUX.1-dev](https://huggingface.co/black-forest-labs/FLUX.1-dev) | 텍스트-이미지 | 12B | 고품질 이미지 생성 |
| HunyuanVideo 1.5 | [hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v](https://huggingface.co/hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v) | 텍스트-투-비디오 | 13B | 고품질 비디오 생성 |

## 설치

### 하드웨어 요구 사항

| 컴포넌트 | 최소 | 권장 |
|-----------|---------|-------------|
| GPU | A100 40GB | A100 80GB / H100 |
| GPU | 4 | 8+ |
| RAM | 128 GB | 256 GB 이상 |
| 저장 | 500 GB SSD | 2 TB NVMe |

pip으로 NeMo Automodel을 설치합니다. 소스에서 설치하는 방법을 포함한 전체 설치 방법은 [NeMo Automodel 설치 가이드](https://docs.nvidia.com/nemo/automodel/latest/guides/installation.html)를 참조하세요.

```bash
pip3 install nemo-automodel
```

대신, 모든 종속성이 포함된 사전 구축된 Docker 컨테이너를 사용할 수 있습니다.

```bash
docker pull nvcr.io/nvidia/nemo-automodel:26.02.00
docker run --gpus all -it --rm --shm-size=8g nvcr.io/nvidia/nemo-automodel:26.02.00
```

> [!WARNING]
> 컨테이너가 종료되면 체크포인트가 손실되므로, 체크포인트 디렉토리를 호스트에 바인드 마운트해야 합니다. 예를 들어, `docker run` 명령에 `-v /host/path/checkpoints:/workspace/checkpoints`를 추가하세요.

## 데이터 준비

NeMo Automodel은 latent space에서 diffusion 모델을 학습합니다. 학습 전에 원본 이미지 또는 비디오를 VAE latents와 text embeddings를 포함하는 `.meta` 파일로 전처리해야 합니다. 이를 통해 각 학습 단계마다 다시 인코딩하는 작업을 방지합니다.

preprocessing 내장 도구를 사용하여 데이터를 인코딩합니다. 이 도구는 사용 가능한 모든 GPU에 작업을 자동으로 분배합니다.

Wan 2.1과 HunyuanVideo 모두에서 비디오 전처리 명령어는 동일하지만 플래그가 다릅니다. Wan 2.1은 `--processor wan`을 `--resolution_preset`과 `--caption_format sidecar`와 함께 사용하는 반면, HunyuanVideo는 프레임 수를 설정하기 위해 `--processor hunyuan`을 `--target_frames`와 함께 사용하고 `--caption_format meta_json`를 사용합니다.

**Wan 2.1:**

```bash
python -m tools.diffusion.preprocessing_multiprocess video \
    --video_dir /data/videos \
    --output_dir /cache \
    --processor wan \
    --resolution_preset 512p \
    --caption_format sidecar
```

**HunyuanVideo:**

```bash
python -m tools.diffusion.preprocessing_multiprocess video \
    --video_dir /data/videos \
    --output_dir /cache \
    --processor hunyuan \
    --target_frames 121 \
    --caption_format meta_json
```

```bash
python -m tools.diffusion.preprocessing_multiprocess image \
    --image_dir /data/images \
    --output_dir /cache \
    --processor flux \
    --resolution_preset 512p
```

### 출력 형식

Preprocessing은 해상도 버킷별로 조직된 캐시 디렉터리를 생성합니다. NeMo Automodel은 버킷 샘플링을 통해 다중 해상도 학습을 지원합니다. 샘플은 공간 해상도별로 그룹화되어 각 배치에 동일한 크기의 샘플만 포함되므로 패딩으로 인한 낭비를 방지합니다.

```
/cache/
├── 512x512/                      # Resolution bucket
│   ├── <hash1>.meta              # VAE latents + text embeddings
│   ├── <hash2>.meta
│   └── ...
├── 832x480/                      # Another resolution bucket
│   └── ...
├── metadata.json                 # Global config (processor, model, total items)
└── metadata_shard_0000.json      # Per-sample metadata (paths, resolutions, captions)
```

> [!TIP]
> 캡션 형식, 입력 데이터 요구 사항 및 사용 가능한 모든 전처리 인수는 [Diffusion Dataset Preparation](https://docs.nvidia.com/nemo/automodel/latest/datasets/diffusion-dataset) 가이드를 참고하세요.

## 학습 구성

Fine-tuning은 두 가지 구성 요소에 의해 구동됩니다:

1. 레시피 스크립트([finetune.py](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/diffusion/finetune/finetune.py))는 학습 루프를 포함하는 Python 진입점입니다. 모델 로딩, dataloader 구성, forward/backward pass 실행, flow matching loss 계산, checkpointing, 로깅을 수행합니다.
2. YAML 구성 파일은 레시피가 사용하는 모든 설정을 지정합니다: 미세 조정할 모델, 데이터 위치, 옵티마이저 하이퍼파라미터, 병렬화 전략 등입니다. 코드를 수정하는 대신 이 파일을 편집하여 학습을 커스터마이징하므로, 1개에서 수백 개의 GPU까지 확장할 수 있습니다.

YAML 필드는 CLI에서도 덮어쓸 수 있습니다:

```bash
torchrun --nproc-per-node=8 examples/diffusion/finetune/finetune.py \
  -c examples/diffusion/finetune/wan2_1_t2v_flow.yaml \
  --optim.learning_rate 1e-5 \
  --step_scheduler.num_epochs 50
```

아래는 Wan 2.1 T2V 1.3B의 fine-tuning을 위한 주석이 달린 config이며, 각 섹션에 대한 설명이 포함되어 있습니다.

```yaml
seed: 42

# ── Experiment tracking (optional) ──────────────────────────────────────────
# Weights & Biases integration for logging metrics, losses, and learning rates.
# Set mode: "disabled" to turn off.
wandb:
  project: wan-t2v-flow-matching
  mode: online
  name: wan2_1_t2v_fm

# ── Model ───────────────────────────────────────────────────────────────────
# pretrained_model_name_or_path: any Hugging Face model ID or local path.
# mode: "finetune" loads pretrained weights; "pretrain" trains from scratch.
model:
  pretrained_model_name_or_path: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
  mode: finetune

# ── Training schedule ───────────────────────────────────────────────────────
# global_batch_size: effective batch across all GPUs.
# Gradient accumulation is computed automatically: global / (local × num_gpus).
step_scheduler:
  global_batch_size: 8
  local_batch_size: 1
  ckpt_every_steps: 1000     # Save a checkpoint every N steps
  num_epochs: 100
  log_every: 2               # Log metrics every N steps

# ── Data ────────────────────────────────────────────────────────────────────
# _target_: the dataloader factory function.
#   Use build_video_multiresolution_dataloader for video models (Wan, HunyuanVideo).
#   Use build_text_to_image_multiresolution_dataloader for image models (FLUX).
# model_type: "wan" or "hunyuan" (selects the correct latent format).
# base_resolution: target resolution for multiresolution bucketing.
data:
  dataloader:
    _target_: nemo_automodel.components.datasets.diffusion.build_video_multiresolution_dataloader
    cache_dir: PATH_TO_YOUR_DATA
    model_type: wan
    base_resolution: [512, 512]
    dynamic_batch_size: false  # When true, adjusts batch per bucket to maintain constant memory
    shuffle: true
    drop_last: false
    num_workers: 0

# ── Optimizer ───────────────────────────────────────────────────────────────
# learning_rate: 5e-6 is a good starting point for fine-tuning.
# Adjust weight_decay and betas for your dataset.
optim:
  learning_rate: 5e-6
  optimizer:
    weight_decay: 0.01
    betas: [0.9, 0.999]

# ── Learning rate scheduler ─────────────────────────────────────────────────
# Supports cosine, linear, and constant schedules.
lr_scheduler:
  lr_decay_style: cosine
  lr_warmup_steps: 0
  min_lr: 1e-6

# ── Flow matching ───────────────────────────────────────────────────────────
# adapter_type: model-specific adapter — must match the model:
#   "simple" for Wan 2.1, "flux" for FLUX.1-dev, "hunyuan" for HunyuanVideo.
# timestep_sampling: "uniform" for Wan, "logit_normal" for FLUX and HunyuanVideo.
# flow_shift: shifts the flow schedule (model-dependent).
# i2v_prob: probability of image-to-video conditioning during training (video models).
flow_matching:
  adapter_type: "simple"
  adapter_kwargs: {}
  timestep_sampling: "uniform"
  logit_mean: 0.0
  logit_std: 1.0
  flow_shift: 3.0
  num_train_timesteps: 1000
  i2v_prob: 0.3
  use_loss_weighting: true

# ── FSDP2 distributed training ──────────────────────────────────────────────
# dp_size: number of GPUs for data parallelism (typically = total GPUs on node).
# tp_size, cp_size, pp_size: tensor, context, and pipeline parallelism.
#   For most fine-tuning, dp_size is all you need; leave others at 1.
fsdp:
  tp_size: 1
  cp_size: 1
  pp_size: 1
  dp_replicate_size: 1
  dp_size: 8

# ── Checkpointing ──────────────────────────────────────────────────────────
# checkpoint_dir: where to save checkpoints (use a persistent path with Docker).
# restore_from: path to resume training from a previous checkpoint.
checkpoint:
  enabled: true
  checkpoint_dir: PATH_TO_YOUR_CKPT_DIR
  model_save_format: torch_save
  save_consolidated: false
  restore_from: null
```

### Config 필드 참조

아래 표에는 최소 요구 설정을 나열했습니다. 모든 모델의 전체 예제 설정은 [NeMo Automodel 예제](https://github.com/NVIDIA-NeMo/Automodel/tree/main/examples/diffusion/finetune)를 참고하세요.

| 섹션 | 필수? | 변경할 사항 |
|---------|-----------|----------------|
| `model` | 예 | `pretrained_model_name_or_path`를 Hugging Face 모델 ID로 설정합니다. `mode: finetune` 또는 `mode: pretrain`을 설정합니다. |
| `step_scheduler` | 예 | `global_batch_size`는 모든 GPU에 걸친 유효 배치 크기입니다. `ckpt_every_steps`는 체크포인트 저장 빈도를 제어합니다. 그라디언트 축적은 자동으로 계산됩니다. |
| `data` | 예 | `cache_dir`를 전처리된 `.meta` 파일이 포함된 경로로 설정하세요. 다른 모델을 사용하는 경우 `_target_`와 `model_type`를 변경하세요. |
| `optim` | 예 | `learning_rate: 5e-6`은 fine-tuning에 좋은 기본값입니다. 데이터셋과 모델에 맞게 조정하세요. |
| `lr_scheduler` | 예 | `lr_decay_style`에는 `cosine`, `linear`, 또는 `constant`를 선택합니다. 점진적인 워밍업을 위해 `lr_warmup_steps`를 설정합니다. |
| `flow_matching` | 예 | `adapter_type`는 모델과 일치해야 합니다(Wan의 경우 `simple`, FLUX의 경우 `flux`, HunyuanVideo의 경우 `hunyuan`). `adapter_kwargs`에 대한 모델별 설정을 참조하세요. |
| `fsdp` | 예 | `dp_size`를 GPU 수로 설정합니다. 멀티 노드 환경에서는 모든 노드에 걸친 GPU 총수로 설정합니다. |
| `checkpoint` | 권장 | Docker 환경에서는 특히 `checkpoint_dir`을 영속적인 경로로 설정하세요. 이전 체크포인트에서 이어서 실행하려면 `restore_from`을 사용하세요. |
| `wandb` | 선택 사항 | Weights & Biases 실험 추적을 활성화하도록 구성합니다. `mode: disabled`를 설정하면 비활성화됩니다. |

## 학습 시작

```bash
torchrun --nproc-per-node=8 \
  examples/diffusion/finetune/finetune.py \
  -c examples/diffusion/finetune/wan2_1_t2v_flow.yaml
```

각 노드에서 다음을 실행하고 `NODE_RANK`을 적절하게 설정하세요:

```bash
export MASTER_ADDR=node0.hostname
export MASTER_PORT=29500
export NODE_RANK=0  # 0 on master, 1 on second node, etc.

torchrun \
  --nnodes=2 \
  --nproc-per-node=8 \
  --node_rank=${NODE_RANK} \
  --rdzv_backend=c10d \
  --rdzv_endpoint=${MASTER_ADDR}:${MASTER_PORT} \
  examples/diffusion/finetune/finetune.py \
  -c examples/diffusion/finetune/wan2_1_t2v_flow_multinode.yaml
```

> [!NOTE]
> 멀티 노드 학습의 경우, YAML에서 `fsdp.dp_size`를 모든 노드에 걸친 **전체** GPU 수로 설정합니다(예: 노드당 8개의 GPU가 있는 2개 노드의 경우 16).

## 생성

After training, fine-tuned checkpoint를 사용하여 텍스트 프롬프트에서 비디오 또는 이미지를 생성합니다.

```bash
python examples/diffusion/generate/generate.py \
  -c examples/diffusion/generate/configs/generate_wan.yaml
```

파인튜닝된 체크포인트를 사용하여:

```bash
python examples/diffusion/generate/generate.py \
  -c examples/diffusion/generate/configs/generate_wan.yaml \
  --model.checkpoint ./checkpoints/step_1000 \
  --inference.prompts '["A dog running on a beach"]'
```

```bash
python examples/diffusion/generate/generate.py \
  -c examples/diffusion/generate/configs/generate_flux.yaml
```

파인튜닝된 체크포인트를 사용하여:

```bash
python examples/diffusion/generate/generate.py \
  -c examples/diffusion/generate/configs/generate_flux.yaml \
  --model.checkpoint ./checkpoints/step_1000 \
  --inference.prompts '["A dog running on a beach"]'
```

```bash
python examples/diffusion/generate/generate.py \
  -c examples/diffusion/generate/configs/generate_hunyuan.yaml
```

파인튜닝된 체크포인트를 사용하여:

```bash
python examples/diffusion/generate/generate.py \
  -c examples/diffusion/generate/configs/generate_hunyuan.yaml \
  --model.checkpoint ./checkpoints/step_1000 \
  --inference.prompts '["A dog running on a beach"]'
```

## Diffusers 통합

NeMo Automodel은 Diffusers를 기반으로 구축되어 모델 로딩 및 추론의 핵심으로 사용합니다. [WanTransformer3DModel](/docs/diffusers/main/ko/api/models/wan_transformer_3d#diffusers.WanTransformer3DModel), [FluxTransformer2DModel](/docs/diffusers/main/ko/api/models/flux_transformer#diffusers.FluxTransformer2DModel), [HunyuanVideoTransformer3DModel](/docs/diffusers/main/ko/api/models/hunyuan_video_transformer_3d#diffusers.HunyuanVideoTransformer3DModel)과 같은 Diffusers 모델 클래스를 사용하여 Hugging Face Hub에서 모델을 직접 로드하고, [WanPipeline](/docs/diffusers/main/ko/api/pipelines/wan#diffusers.WanPipeline) 및 [FluxPipeline](/docs/diffusers/main/ko/api/pipelines/flux#diffusers.FluxPipeline)과 같은 Diffusers 파이프라인을 통해 출력을 생성합니다.

이 통합은 Diffusers 사용자에게 다음과 같은 여러 이점을 제공합니다:

- **체크포인트 변환 없음**: Hub의 사전 학습된 가중치는 바로 사용할 수 있습니다. `pretrained_model_name_or_path`를 Diffusers 형식의 모델 ID로 지정하고 즉시 학습을 시작할 수 있습니다.
- **Day-0 모델 지원**: 새로운 diffusion 모델이 Diffusers에 추가되고 Hub에 업로드되면, 전용 학습 스크립트를 기다릴 필요 없이 NeMo Automodel로 fine-tuning할 수 있습니다.
- **Pipeline 호환 출력**: fine-tuned 체크포인트는 추론을 위해 Diffusers 파이프라인에 직접 다시 로드하거나, Hub에 공유하거나, 양자화 및 컴파일과 같은 도구를 사용하여 추가 최적화할 수 있는 형식으로 저장됩니다.
- **Diffusers 모델의 확장 가능한 학습**: NeMo Automodel은 내장된 Diffusers 학습 스크립트가 제공하는 범위를 넘어선 분산 학습 기능(FSDP2, 멀티노드, 멀티레졸루션 버킷팅)을 추가하되, 동일한 모델 및 파이프라인 인터페이스를 유지합니다.
- **공유 생태계**: Diffusers 생태계의 모든 모델, LoRA 어댑터, 파이프라인 구성 요소는 학습 및 추론 워크플로우 전반에 걸쳐 호환성을 유지합니다.

## NVIDIA 팀

- Pranav Prashant Thombre, pthombre@nvidia.com
- Linnan Wang, linnanw@nvidia.com
- Alexandros Koumparoulis, akoumparouli@nvidia.com

## 리소스

- [NeMo Automodel GitHub](https://github.com/NVIDIA-NeMo/Automodel)
- [Diffusion Fine-Tuning 가이드](https://docs.nvidia.com/nemo/automodel/latest/recipes-e2e-examples/diffusion-fine-tuning)
- [Diffusion 데이터셋 준비](https://docs.nvidia.com/nemo/automodel/latest/datasets/diffusion-dataset)
- [확산 모델 지원 범위](https://docs.nvidia.com/nemo/automodel/latest/model-coverage/diffusion.html)
- [Transformers용 NeMo Automodel (LLM/VLM 파인튜닝)](https://huggingface.co/docs/transformers/en/community_integrations/nemo_automodel_finetuning)