En este tutorial, implementamos una implementación práctica y avanzada de NVIDIA Transformer Engine en Python, centrándonos en cómo se puede explorar la aceleración de precisión mixta en un flujo de trabajo de aprendizaje profundo realista. Configuramos el entorno, verificamos la preparación de GPU y CUDA, intentamos instalar los componentes necesarios de Transformer Engine y manejamos los problemas de compatibilidad con elegancia para que la computadora portátil siga siendo ejecutable incluso cuando no se puede crear la extensión completa. A medida que avanzamos en cada paso, construimos redes de profesores y estudiantes, comparamos una ruta de referencia de PyTorch con una ruta habilitada para Transformer Engine, entrenamos ambos modelos, comparamos su velocidad y uso de memoria y visualizamos los resultados, lo que nos brinda una comprensión práctica clara de cómo se estructuran en la práctica los flujos de trabajo de capacitación orientados al rendimiento.
importar sistema operativo importar sistema importar json importar tiempo importar matemáticas importar aleatorio importar Shuil importar plataforma importar subproceso importar estadísticas def ejecutar(cmd, check=True): print("n[RUN]", " ".join(cmd)) resultado = subprocess.run(cmd, text=True, capture_output=True) if result.stdout.strip(): print(result.stdout[-4000:]) if result.returncode != 0 y result.stderr.strip(): print(result.stderr[-4000:]) if check y result.returncode != 0: levanta subprocess.CalledProcessError(result.returncode, cmd) devuelve resultado def has_cmd(name): devuelveshutil.what(name) no es Ninguno run([sys.executable, "-m", "pip", "install", "-q", "–upgrade", "pip"]) run([sys.executable, "-m", "pip", "install", "-q", "ninja", "packaging", "matplotlib"]) import torch import torch.nn as nn import torch.nn.function as F import matplotlib.pyplot as plt afirmar torch.cuda.is_available(), "Este portátil necesita un tiempo de ejecución de GPU en Colab". gpu_name = torch.cuda.get_device_name(0) cc_major, cc_minor = torch.cuda.get_device_capability(0) cuda_runtime = torch.version.cuda python_version = sys.version.split()[0]torch_version = torch.__version__ cuda_home = os.environ.get("CUDA_HOME", "/usr/local/cuda") nvcc_path =shutil. Which("nvcc") o os.path.join(cuda_home, "bin", "nvcc") cudnn_header_candidates = [ os.path.join(cuda_home, "include", "cudnn.h"), "/usr/include/cudnn.h", "/usr/local/include/cudnn.h", ] nvcc_exists = os.path.exists(nvcc_path) cudnn_header_exists = any(os.path.exists(p) para p en cudnn_header_candidates) print("=" * 120) print("COMPROBACIÓN DEL ENTORNO") print("=" * 120) print(json.dumps({ "python": python_version, "platform": platform.platform(), "torch": torch_version, "torch_cuda": cuda_runtime, "gpu_name": gpu_name, "compute_capability": f"{cc_major}.{cc_minor}", "cuda_home": cuda_home, "nvcc_exists": nvcc_exists, "nvcc_path": nvcc_path si nvcc_exists else Ninguno, "cudnn_header_exists": cudnn_header_exists, }, sangría=2)) print("=" * 120)
Preparamos el entorno Colab importando las bibliotecas Python requeridas, definiendo una función auxiliar para ejecutar comandos de shell e instalando las dependencias principales para el tutorial. Luego importamos PyTorch y Matplotlib, verificamos que haya una GPU disponible y recopilamos detalles clave del entorno, incluido el nombre de la GPU, la versión de CUDA, la versión de Python y las rutas del kit de herramientas. Esto nos brinda una visión clara del estado del sistema antes de intentar cualquier instalación de Transformer Engine o ejecución del modelo.
te_available = False te_mode = "fallback" te_import_error = Ninguno try: run([sys.executable, "-m", "pip", "install", "-q", "transformer_engine[core_cu12]"]) excepto Excepción como e: print("Falló la instalación de la rueda central:", repr(e)) can_try_te_torch = nvcc_exists y cudnn_header_exists if can_try_te_torch: env = os.environ.copy() env["NVTE_FRAMEWORK"] = "pytorch" env["MAX_JOBS"] = "1" env["NVTE_BUILD_THREADS_PER_JOB"] = "1" env["CUDA_PATH"] = cuda_home env["CUDA_HOME"] = cuda_home intente: print("nIntentando construir la extensión PyTorch para Transformer Engine…") result = subprocess.run( [sys.executable, "-m", "pip", "install", "-q", "–no-build-isolation", "transformer_engine[pytorch]"], text=True, capture_output=True, env=env, ) if result.stdout.strip(): print(result.stdout[-4000:]) if result.returncode != 0 y result.stderr.strip(): print(result.stderr[-4000:]) if result.returncode == 0: importar transformador_engine.pytorch como te de transformador_engine.common receta de importación te_available = True te_mode = "transformer_engine" else: te_import_error = result.stderr[-4000:] if result.stderr else "Error de compilación de pip desconocido" excepto Excepción como e: te_import_error = repr(e) else: te_import_error = "Faltan encabezados nvcc o cuDNN en este tiempo de ejecución de Colab, por lo que la extensión TE PyTorch no se puede crear aquí". si te_available: intente: fp8_available, fp8_reason = te.is_fp8_available(return_reason=True) excepto excepción como e: fp8_available, fp8_reason = False, f"No se pudo consultar la disponibilidad de FP8: {e}" intente: bf16_available = te.is_bf16_available() excepto excepción: bf16_available = torch.cuda.is_bf16_supported() else: fp8_available = False fp8_reason = "Transformer Engine no está instalado; usando la ruta alternativa de PyTorch." bf16_available = torch.cuda.is_bf16_supported() amp_dtype = torch.bfloat16 if bf16_available else torch.float16 print("n" + "=" * 120) print("ESTADO DE INSTALACIÓN") print("=" * 120) print(json.dumps({ "te_available": te_available, "te_mode": te_mode, "fp8_available": fp8_available, "fp8_reason": fp8_reason, "te_import_error": te_import_error, "amp_dtype": str(amp_dtype), }, indent=2)) print("=" * 120) dispositivo = "cuda" random.seed(42) torch.manual_seed(42) torch.cuda.manual_seed_all(42) if te_available: fp8_recipe = receta.DelayedScaling(margin=0, fp8_format=recipe.Format.E4M3) def baseline_autocast(): devuelve torch.autocast(device_type="cuda", dtype=amp_dtype) def te_forward_context(use_fp8): si te_available y use_fp8: devuelve te.autocast(enabled=True, receta=fp8_recipe) devuelve baseline_autocast()
Intentamos instalar el paquete principal de Transformer Engine y luego verificamos si el tiempo de ejecución de Colab puede crear la extensión PyTorch verificando la presencia de encabezados nvcc y cuDNN. Si el entorno lo admite, intentamos instalar el backend de Transformer Engine PyTorch y luego inspeccionamos si FP8 y BF16 están disponibles en el hardware actual. También configuramos el modo de precisión y definimos los contextos de autocast que luego nos permiten cambiar entre precisión mixta estándar y ejecución de Transformer Engine.
clase TeacherNet(nn.Module): def __init__(self, tamaño_oculto=512, tamaño_intermedio=2048, num_layers=3, tamaño_vocab=4096): super().__init__() self.embed = nn.Embedding(tamaño_vocab, tamaño_oculto) self.layers = nn.ModuleList([ nn.Sequential( nn.LayerNorm(tamaño_oculto), nn.Linear(tamaño_oculto, tamaño_intermedio), nn.GELU(), nn.Linear(tamaño_intermedio, tamaño_oculto), ) para _ en rango(núm_capas) ]) self.head = nn.Linear(tamaño_oculto, tamaño_oculto) def forward(self, token_ids): x = self.embed(token_ids) para la capa en self.layers: x = x + Layer(x) return self.head(x) clase BaselineStudent(nn.Module): def __init__(self, Hidden_size=512, middle_size=2048, num_layers=3, vocab_size=4096): super().__init__() self.embed = nn.Embedding(vocab_size, Hidden_size) self.norms = nn.ModuleList([nn.LayerNorm(hidden_size) for _ in range(num_layers)]) self.fc1 = nn.ModuleList([nn.Linear(hidden_size, intermedio_size) for _ in range(num_layers)]) self.fc2 = nn.ModuleList([nn.Linear(intermediate_size, Hidden_size) for _ in range(num_layers)]) self.head = nn.Linear(hidden_size, Hidden_size) def forward(self, token_ids): x = self.embed(token_ids) for ln, fc1, fc2 in zip(self.norms, self.fc1, self.fc2): residual = x x = ln(x) x = fc1(x) x = F.gelu(x, aproximado="tanh") x = fc2(x) x = x + retorno residual self.head(x) if te_available: clase TEStudent(nn.Module): def __init__(self, Hidden_size=512, middle_size=2048, num_layers=3, vocab_size=4096): super().__init__() self.embed = nn.Embedding(vocab_size, Hidden_size) self.norms = nn.ModuleList([te.LayerNorm(hidden_size) for _ in range(num_layers)]) self.fc1 = nn.ModuleList([te.Linear(hidden_size, intermedio_size, sesgo=True) for _ in range(num_layers)]) self.fc2 = nn.ModuleList([te.Linear(intermediate_size, Hidden_size, sesgo=True) para _ en rango(num_layers)]) self.head = te.Linear(hidden_size, Hidden_size, sesgo=True) def forward(self, token_ids, use_fp8=False): x = self.embed(token_ids) con te_forward_context(use_fp8): para ln, fc1, fc2 en zip(self.norms, self.fc1, self.fc2): residual = x x = ln(x) x = fc1(x) x = F.gelu(x, aproximado="tanh") x = fc2(x) x = x + residual x = self.head(x) return x else: clase TEStudent(nn.Module): def __init__(self, hide_size=512, middle_size=2048, num_layers=3, tamaño_vocab=4096): super().__init__() self.embed = nn.Embedding(tamaño_vocab, tamaño_oculto) self.norms = nn.ModuleList([nn.LayerNorm(tamaño_oculto) para _ en rango(núm_capas)]) self.fc1 = nn.ModuleList([nn.Linear(tamaño_oculto, tamaño_intermedio) para _ en rango(núm_capas)]) self.fc2 = nn.ModuleList([nn.Linear(tamaño_intermedio, tamaño_oculto) para _ en rango(núm_capas)]) self.head = nn.Linear(tamaño_oculto, tamaño_oculto) def forward(self, token_ids, use_fp8=False): x = self.embed(token_ids) con baseline_autocast(): para ln, fc1, fc2 en zip(self.norms, self.fc1, self.fc2): residual = x x = ln(x) x = fc1(x) x = F.gelu(x, aproximado="tanh") x = fc2(x) x = x + residual x = self.head(x) return x def count_params(model): return sum(p.numel() for p in model.parameters() if p.requires_grad) def format_millions(n): devolver f"{n / 1e6:.2f}M"
Definimos las arquitecturas de redes neuronales utilizadas a lo largo del tutorial, incluido el modelo de maestro, el modelo de estudiante de referencia y la ruta de estudiante de Transformer Engine. Mantenemos las estructuras del modelo alineadas para que la comparación siga siendo significativa y al mismo tiempo permitimos que la ruta TE intercambie capas de Transformer Engine cuando la extensión esté disponible. También definimos pequeñas funciones de utilidad para contar parámetros y formatear el tamaño del modelo, que nos ayudan a inspeccionar la escala de los modelos antes de que comience el entrenamiento.
tamaño_oculto = 512 tamaño_intermedio = 2048 num_layers = 3 tamaño_vocab = 4096 seq_len = 128 tamaño_lote = 8 pasos = 25 benchmark_iters = 20 lr = 2e-4 peso_decay = 1e-2 profesor = TeacherNet(tamaño_oculto, tamaño_intermedio, num_capas, tamaño_vocab).to(dispositivo).eval() baseline_model = BaselineStudent(tamaño_oculto, tamaño_intermedio, núm_capas, tamaño_vocab).to(dispositivo) te_model = TEStudent(tamaño_oculto, tamaño_intermedio, núm_capas, tamaño_vocab).to(dispositivo) optimizador_baseline = torch.optim.AdamW(baseline_model.parameters(), lr=lr, peso_decay=peso_decay) optimizador_te = torch.optim.AdamW(te_model.parameters(), lr=lr, Weight_decay=weight_decay) print("Parámetros del profesor:", format_millions(count_params(profesor))) print("Parámetros de referencia:", format_millions(count_params(baseline_model))) print("Parámetros de ruta TE:", format_millions(count_params(te_model))) def make_batch(batch_size, seq_len, vocab_size, dispositivo): tokens = torch.randint(0, vocab_size, (batch_size, seq_len), dispositivo=dispositivo) con torch.no_grad(): target = profesor(tokens) devolver tokens, target def peak_mem_mb(): return torch.cuda.max_memory_allocated() / (1024 ** 2) def train_baseline_step(): baseline_model.train() optimizador_baseline.zero_grad(set_to_none=True) tokens, target = make_batch(batch_size, seq_len, vocab_size, dispositivo) con baseline_autocast(): pred = baseline_model(tokens) loss = F.mse_loss(pred, target) loss.backward() optimizador_baseline.step() devuelve float(loss.detach().item()) def train_te_step(use_fp8): te_model.train() optimizador_te.zero_grad(set_to_none=True) tokens, destino = make_batch(batch_size, seq_len, vocab_size, dispositivo) pred = te_model(tokens, use_fp8=use_fp8) pérdida = F.mse_loss(pred, target) pérdida.backward() optimizador_te.step() return float(loss.detach().item())
Configuramos los hiperparámetros del experimento principal, creamos instancias de todos los modelos en la GPU y creamos los optimizadores que se utilizarán durante el entrenamiento. También imprimimos los recuentos de parámetros para confirmar que las rutas de referencia y TE son comparables en términos de tamaño del modelo. Además, definimos la lógica de generación por lotes, la función de seguimiento de la memoria y las funciones de pasos de entrenamiento individuales que ejecutan un paso de optimización para cada ruta del modelo.
pérdidas_linea_base =[]te_pérdidas =[]mode_name = "TE-FP8" if (te_available y fp8_available) else ("TE-BF16/FP16" if te_available else "Fallback-PyTorch") print("n" + "=" * 120) print("ENTRENAMIENTO") print("=" * 120) para el paso en el rango(1, pasos + 1): b_loss = train_baseline_step() t_loss = train_te_step(use_fp8=fp8_available) baseline_losses.append(b_loss) te_losses.append(t_loss) si paso == 1 o paso % 5 == 0 o paso == pasos: print(f"step={step:02d} | baseline_loss={b_loss:.6f} | te_path_loss={t_loss:.6f} | mode={mode_name}") @torch.no_grad() def evalua_model(modelo, is_te=False, use_fp8=False, eval_batches=8): model.eval() vals =[]para _ en rango(eval_batches): tokens, target = make_batch(batch_size, seq_len, vocab_size, dispositivo) if is_te: pred = model(tokens, use_fp8=use_fp8) else: with baseline_autocast(): pred = model(tokens) vals.append(float(F.mse_loss(pred, target).item())) return suma(vals) / len(vals) baseline_eval = evaluar_modelo(baseline_model, is_te=False) te_eval = evaluar_model(te_model, is_te=True, use_fp8=fp8_available) def benchmark_train_step(modelo, optimizador, is_te=False, use_fp8=False, calentamiento=5, iters=20): times_ms =[]mems_mb=[]para _ en rango (calentamiento): tokens optimizadores.zero_grad(set_to_none=True), destino = make_batch(tamaño_batch, seq_len, tamaño vocab, dispositivo) if is_te: pred = modelo(tokens, use_fp8=use_fp8) else: con baseline_autocast(): pred = modelo(tokens) pérdida = F.mse_loss(pred, objetivo) loss.backward() optimizador.step() torch.cuda.synchronize() para _ en rango(iters): torch.cuda.reset_peak_memory_stats() optimizador.zero_grad(set_to_none=True) tokens, target = make_batch(batch_size, seq_len, vocab_size, dispositivo) start = time.perf_counter() if is_te: pred = model(tokens, use_fp8=use_fp8) else: con baseline_autocast(): pred = model(tokens) loss = F.mse_loss(pred, target) loss.backward() optimizador.step() torch.cuda.synchronize() end = time.perf_counter() times_ms.append((end – start) * 1000.0) mems_mb.append(peak_mem_mb()) return { "mean_ms": stats.mean(times_ms), "median_ms": stats.median(times_ms), "max_memory_mb": max(mems_mb), } baseline_bench = benchmark_train_step(baseline_model, optimizador_baseline, is_te=False, use_fp8=False, iters=benchmark_iters) te_bench = benchmark_train_step(te_model, optimizador_te, is_te=True, use_fp8=fp8_available, iters=benchmark_iters)
Ejecutamos el ciclo de entrenamiento principal tanto para el modelo de referencia como para la ruta TE, rastreando sus pérdidas en múltiples pasos. Luego definimos y ejecutamos la función de evaluación para medir qué tan bien coincide cada modelo con los resultados del maestro después de la capacitación. Finalmente, implementamos la rutina de evaluación comparativa para medir el tiempo de ejecución por paso y el uso máximo de memoria CUDA, lo que permite una comparación cuantitativa de las características de rendimiento.
resumen = { "gpu_name": gpu_name, "compute_capability": f"{cc_major}.{cc_minor}", "te_available": te_available, "fp8_available": fp8_available, "fp8_reason": fp8_reason, "mode": mode_name, "baseline_eval_mse": baseline_eval, "te_path_eval_mse": te_eval, "baseline_mean_step_ms": baseline_bench["mean_ms"], "te_path_mean_step_ms": te_bench["mean_ms"], "baseline_peak_mem_mb": baseline_bench["max_memory_mb"], "te_path_peak_mem_mb": te_bench["max_memory_mb"], } print("n" + "=" * 120) print("RESUMEN") print("=" * 120) print(json.dumps(summary, indent=2)) plt.figure(figsize=(10, 5)) plt.plot(baseline_losses, label="Pérdida de referencia") plt.plot(te_losses, label=f"{mode_name} loss") plt.xlabel("Paso de entrenamiento") plt.ylabel("Pérdida de MSE") plt.title("Comparación de pérdidas de entrenamiento") plt.legend() plt.grid(True) plt.show() plt.figure(figsize=(8, 5)) plt.bar(["Baseline", mode_name], [baseline_bench["mean_ms"], te_bench["mean_ms"]]) plt.ylabel("Tiempo medio de paso del tren (ms)") plt.title("Comparación de velocidades") plt.grid(True, axis="y") plt.show() plt.figure(figsize=(8, 5)) plt.bar(["Baseline", mode_name], [baseline_bench["max_memory_mb"], te_bench["max_memory_mb"]]) plt.ylabel("Memoria máxima (MB)") plt.title("Comparación de memoria CUDA máxima") plt.grid(True, axis="y") plt.show()
Reunimos todas las métricas finales en un diccionario resumido e imprimimos los resultados consolidados del experimento en un formato estructurado. Luego generamos visualizaciones de la pérdida de entrenamiento, el tiempo medio de los pasos de entrenamiento y el uso máximo de la memoria para interpretar de manera más intuitiva las diferencias entre las rutas de referencia y TE. Esta sección final nos ayuda a pasar de números brutos a conocimientos prácticos al mostrar cómo se comportan las dos implementaciones en términos de precisión, velocidad y memoria.
En conclusión, creamos mucho más que un simple tutorial de instalación; Creamos un proceso experimental completo que nos ayuda a comprender cómo encaja NVIDIA Transformer Engine en el entrenamiento de modelos modernos acelerados por GPU. Probamos el entorno de ejecución, lo adaptamos a las limitaciones de Colab, conservamos una ruta alternativa funcional y luego entrenamos, evaluamos y comparamos dos implementaciones una al lado de la otra para observar diferencias prácticas en eficiencia, comportamiento de precisión y uso de recursos. Al final, entendimos cómo usar Transformer Engine en un entorno compatible con Colab y obtuvimos una base reutilizable que podemos extender a arquitecturas de transformadores más grandes, escenarios de evaluación comparativa más completos y flujos de trabajo de optimización más orientados a la producción.
Consulte los códigos completos/cuaderno aquí. Además, no dude en seguirnos en Twitter y no olvide unirse a nuestro SubReddit de más de 120.000 ML y suscribirse a nuestro boletín. ¡Esperar! estas en telegrama? Ahora también puedes unirte a nosotros en Telegram.
¿Necesita asociarse con nosotros para promocionar su repositorio de GitHub O su página principal de Hugging O su lanzamiento de producto O seminario web, etc.? Conéctate con nosotros
