Kubernetes 调度 RTX 4090D 节点部署 PyTorch 2.8 运行 AIGC 任务

持久化存储配置

Pod 销毁后数据不能丢，所以得配 PV 和 PVC。创建一个 storage.yaml 文件来定义存储卷：

apiVersion: v1
kind: PersistentVolume
metadata:
  name: gpu-data-pv
spec:
  capacity:
    storage: 40Gi
  accessModes:
    - ReadWriteOnce
  hostPath:
    path: /mnt/data
---
apiVersion: v1
kind: PersistentVolumeClaim
metadata:
  name: gpu-data-pvc
spec:
  accessModes:
    - ReadWriteOnce
  resources:
    requests:
      storage: 40Gi

然后在 Pod 配置里挂载这个卷，把 /data 映射到 PVC：

# 在原有 pod 配置中添加 volumes 部分
volumes:
- name: gpu-data
  persistentVolumeClaim:
    claimName: gpu-data-pvc
containers:
- name: pytorch-container
  volumeMounts:
  - mountPath: /data
    name: gpu-data

运行 AIGC 任务

大模型推理示例

文本生成是基础场景。将以下代码保存为 /workspace/text_generation.py，记得替换实际的模型路径。

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_path = "/workspace/models/your-model"  # 替换为实际模型路径
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    torch_dtype=torch.float16,
    device_map="auto"
)

input_text = "人工智能的未来发展将会"
inputs = tokenizer(input_text, return_tensors="pt").to("cuda")
outputs = model.generate(**inputs, max_new_tokens=50)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

视频生成任务示例

如果是视频生成，Diffusers 库通常更合适。注意显存占用较大，建议开启 fp16 精度：

from diffusers import DiffusionPipeline
import torch

pipeline = DiffusionPipeline.from_pretrained(
    "your-video-model",
    torch_dtype=torch.float16,
    variant="fp16"
).to("cuda")

prompt = "A futuristic city at night, cyberpunk style"
video_frames = pipeline(prompt, num_frames=24).frames
video_frames[0].save("/workspace/output/video.gif", save_all=True, append_images=video_frames[1:], duration=100, loop=0)

性能优化建议

显存优化技巧

RTX 4090D 虽然拥有 24GB 显存，但跑大模型时依然紧张。试试 4bit 量化，能显著降低内存需求：

from transformers import BitsAndBytesConfig

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.float16
)

model = AutoModelForCausalLM.from_pretrained(
    model_path,
    quantization_config=bnb_config,
    device_map="auto"
)

使用 xFormers 加速

对于扩散模型，启用 xFormers 往往能带来明显的速度提升：

# 在 Diffusers 管道中启用
pipeline.enable_xformers_memory_efficient_attention()

常见问题解决

GPU 无法识别

如果 Pod 里看不到卡，按顺序排查：

1. 节点是否正确打标（kubectl get nodes --show-labels）
2. NVIDIA Device Plugin 是否安装正常
3. 驱动版本是否匹配（推荐 550.90.07 及以上）

显存不足

遇到 OOM 错误怎么办？

1. 优先尝试模型量化（4bit/8bit）
2. 减小 batch size
3. 开启梯度检查点（Gradient Checkpointing）

首次加载缓慢

大模型第一次加载通常需要 1-3 分钟，这是磁盘 IO 和权重解压的正常过程。如果频繁复用，可以考虑预加载模型或使用更快的 NVMe 存储介质。

总结

这套方案的核心在于利用 Kubernetes 的资源调度能力，结合 RTX 4090D 的高性价比算力，实现 AIGC 任务的自动化部署。关键点包括正确配置节点标签、合理声明 GPU 资源、持久化存储模型数据，以及针对显存进行必要的量化优化。按照上述步骤操作，你可以高效地在集群中运行各类生成式 AI 任务。