La parte 2 se cerraba con una firma: un batch de secuencias de filas de embedding es un tensor de forma [B, T, d]. En todo este stack no hay nada más parecido a un tipo que esa línea, así que aquí la leo como se lee la firma de una función y luego enseño qué hace el runtime cuando la firma está mal. Lo habitual es que no haga nada.
He leído la regla de broadcasting tal como la escribe torch y después he escrito a propósito cuatro bugs de forma para ver cuáles detectaba el runtime. Solo ha detectado uno.
Una forma es una firma
Un tensor es un array de N dimensiones junto con una descripción de cómo está colocado en memoria. [B, T, d] dice que hay B secuencias en el batch, T posiciones en cada una y d números en cada posición. Léelo como si fuera f(batch: List[List[Vec]]), con una diferencia que importará más adelante, y es que el runtime solo mira los tamaños y los nombres le dan igual.
La descripción de la memoria es lo que un ingeniero de software ya conoce. Un tensor de forma [4, 8, 16] tiene strides (128, 16, 1), así que dar un paso por el primer eje supone saltarse 128 números. Si transpones los dos primeros ejes, la forma pasa a [8, 4, 16] y los strides a (16, 128, 1). Es la misma memoria leída en otro orden, y data_ptr no cambia. Por eso transpose sale gratis, un view que aplane el resultado se niega y lo dice en el mensaje de error, y reshape copia cuando no le queda otro remedio. Un tensor es un puntero, una forma y una tupla de strides. Las operaciones de vista que vas a usar son aritmética sobre las dos últimas.
Cada operación tiene también su propia firma de formas, que dice qué entra, qué sale y qué eje desaparece. mean(dim=1) sobre [B, T, d] devuelve [B, d], ya sin la T. Una capa lineal lleva [B, T, d] a [B, T, d_out] y solo cambia el último eje. Escribir la firma al lado de la llamada es la disciplina de la que trata esta entrada, y si no la escribes tú, no la va a escribir nadie.
La regla cabe en cinco líneas
Dos tensores de formas distintas se pueden sumar igualmente, y cómo se suman lo decide una regla llamada broadcasting. torch tiene una copia legible en Python, torch.broadcast_shapes, que llama a _broadcast_shapes en torch/_refs/__init__.py (líneas 388 a 470 en la v2.13.0). Quitando la maquinaria de formas simbólicas, la regla queda así:
common_shape = [1] * reduce(max, (len(shape) for shape in shapes))for arg_idx, shape in enumerate(shapes): for idx in range(-1, -1 - len(shape), -1): # align from the right if common_shape[idx] == 1: common_shape[idx] = shape[idx] # 1 stretches to anything if shape[idx] == 1: continue # ... in either direction torch._check(common_shape[idx] == shape[idx], # otherwise sizes must match lambda: f"Attempting to broadcast a dimension of length {shape[idx]} at {idx}!")return common_shapeLas formas se alinean por la derecha, un tamaño 1 se estira hasta igualar al otro, los tamaños iguales pasan y cualquier otra combinación lanza un error. No hay más. En modo eager, la misma regla se ejecuta en C++ y falla con un mensaje más corto, que es el que muestra la demo de abajo. La regla existe para que puedas sumar un sesgo de forma [d] a [B, T, d] sin bucles ni copias.
El resto de la entrada gira en torno a una propiedad de la regla: decide por el tamaño y es ciega al significado. [8, 16] + [16] funciona porque 16 es igual a 16, y [8, 16] + [8] lanza un error porque 8 no es 16 ni es 1. La regla no tiene ni idea de si el 16 de la derecha es una longitud de secuencia, una dimensión oculta o un batch, y cuando ese eje sea el equivocado, tú tampoco lo sabrás. Y si dos ejes coinciden en tamaño por casualidad, la regla no tiene manera de distinguirlos. El científico de datos lleva años apoyándose en ella con numpy, y el ingeniero de software se la encuentra como una regla de coerción sin tipos declarados, que es como deberían leerla los dos.
El bug que solo aparece con un batch de dieciséis
En el primer bug intervienen las puntuaciones de atención, de forma [B, T, T] (una matriz T × T por secuencia), y una máscara de padding de forma [B, T] que dice qué posiciones de cada secuencia son reales. En su versión aditiva (0 en las reales y -inf en el padding), basta con sumarla a las puntuaciones para apagar el padding. Lo correcto es scores + mask[:, None, :]. Con el None, la máscara pasa a ser [B, 1, T], se estira por el eje de las queries y enmascara las keys.
Ahora quita el None. Con B = 8 y T = 16, scores + mask lanza un error porque 16 no es 8, pero con B = 16 y T = 16 se ejecuta. Con B = 1 también se ejecuta y además acierta por casualidad, porque 1 es justo lo que habría insertado el None.
La regla alinea [16, 16] con [16, 16, 16] por la derecha y lee la máscara como [1, B, T], de modo que la máscara pensada para la secuencia b cae en la fila de query b de todas las secuencias. Con longitudes aleatorias he medido un tercio de las posiciones mal enmascaradas y una quinta parte del peso de atención sobre el padding. La softmax que viene detrás da valores finitos, la pérdida también y la ejecución sigue adelante. No salta ningún aviso, porque no hay nada de qué avisar. La regla se ha cumplido.
[B, T, T] más una máscara de padding. Arrastra los dos tamaños y elige la forma de la máscara. El veredicto es el de la regla de broadcasting de torch, y lo único que mira es si los tamaños coinciden.Mueve los dos tamaños hasta que coincidan y verás cómo el veredicto pasa de un error a un resultado silencioso. Después elige [B, T, 1], una máscara que se ejecuta con cualquier tamaño y siempre está mal, porque apaga las queries cuando debería apagar las keys. Esta sí se nota, ya que la fila de una query de padding es toda -inf y su softmax da NaN. En los tres casos la regla es la misma.
El segundo bug es el clásico: predicciones de forma [256] y valores objetivo de forma [256, 1] que se restan para calcular un error cuadrático. La regla estira los dos a [256, 256] y enfrenta cada predicción con cada objetivo. La media de todo eso sigue siendo un número, y con él se entrena.
Con este bug he entrenado una regresión lineal de una sola variable durante dos mil pasos. La pérdida ha bajado hasta estabilizarse en 9,7. La pendiente ajustada ha salido 0,000 y la ordenada en el origen igual a la media de los objetivos, con un R² de 0,0. El modelo ha aprendido a ignorar la entrada, que es lo mejor que puede hacer cuando la pérdida lo compara con los objetivos de todos los demás. La pérdida de la biblioteca, F.mse_loss, avisa de este desajuste. La escrita a mano, ((pred - target) ** 2).mean(), no avisa, y es la que acaba en los notebooks.
El tercer bug nunca da error en el sitio donde ocurre. Sobre un tensor de forma [16, 16, 64], tanto mean(dim=0) como mean(dim=1) devuelven [16, 64]. Los números son distintos, pero con la misma forma nada de lo que venga después puede notarlo. Sobre [3, 5, 7], esas mismas dos llamadas devuelven [5, 7] y [3, 7], y lo primero que intente combinarlas lanza un error. El bug estaba en los dos casos y solo los tamaños han decidido si se veía.
El cuarto hace de control. [3, 5, 7] más [3, 5] falla porque 7 no es 5, y si se ha detectado es por lo único que les faltaba a los otros tres, que los tamaños fueran distintos.
Para ser justos con la regla, hay que decir que hace lo que se le pide. Alternativa hay: un runtime que lleve los nombres de los ejes y se niegue a sumar T con B. torch publicó como prototipo los named tensors, que nunca pasaron de esa fase, y la 2.13 los ha eliminado para quitar sobrecarga y aligerar el código. Si no cuajaron fue en parte porque todas las bibliotecas del stack habrían tenido que ponerse de acuerdo en los nombres, y lo único que comparten todas son los tamaños. La regla es correcta, solo que no es una comprobación de tipos, y en todo el stack no hay nada que lo sea.
El padding es una decisión de forma
Las secuencias no vienen todas con la misma longitud. Tres secuencias de longitudes 5, 3 y 1 solo forman un tensor [3, 5] si se rellenan, y con esas longitudes el 40 % del tensor es padding. A partir de ahí, toda operación sobre T necesita saber qué celdas son reales, y ese dato es un segundo tensor, la máscara, que tiene que viajar con el primero.
El lado por el que se rellena forma parte de la firma. Con padding por la derecha, el último token real está en el índice length - 1, que para las tres de arriba es 4, 2 y 0. Con padding por la izquierda está en el 4 para las tres. Los decodificadores que generan token a token prefieren esta segunda opción, porque la posición que hay que leer es siempre la última. Las dos funcionan y el bug está en mezclarlas. El código pensado para padding por la derecha, si recibe un batch rellenado por la izquierda, lee padding, o el token equivocado, donde espera el último.
La máscara no es opcional ni siquiera cuando el modelo no usa atención. Si haces pooling de media de ese batch [3, 5] sin ella, el vector de la secuencia de longitud 1 promedia una fila real con cuatro de padding. Con arange(1, 31).reshape(3, 5, 2), la media de esa secuencia sin máscara es [25, 26], y con máscara es [21, 22], que es su única fila real. Por eso el código de pooling de la parte 2 divide por la suma de la máscara, y esa división es la firma puesta por escrito.
La memoria es forma por bytes
Lo que cuesta un tensor es pura aritmética. Unas activaciones de forma [32, 2048, 4096] son 268 millones de números, que ocupan 1,07 GB en fp32, 537 MB en bf16 y 268 MB en int8. Esa es la factura, y para calcularla no hace falta saber nada del modelo, basta con la forma y el ancho de un número. De esa cifra depende que un batch quepa o no en la tarjeta.
T × T completa. Cada cabeza mide 64 de ancho, así que hay d / 64 cabezas. Arrastra los tres tamaños y cambia el formato numérico: los bytes de cada fila son la forma por el ancho de un número, y los ids de token son int64 sea cual sea el formato. En int8 la cifra solo pone precio a la forma, porque en una ejecución real las puntuaciones y los pesos no se guardan en int8. Las dos últimas filas llevan T dos veces.Mueve los tres tamaños y cambia el formato. Con los valores iniciales, un bloque cuesta 77 GB en fp32 y las dos filas de T × T suponen el 89 % del total, pero si bajas T a 128 se quedan en un tercio. De esas dos filas trata la parte siguiente.
Lo que sacrifica cada formato también es aritmética. Donde fp32 guarda 23 bits de mantisa, bf16 guarda 7 y conserva todo el exponente de fp32. Con esos 7 bits escribe π como 3,140625 y redondea 1 + 0.004 a 1.0078125, pero su rango es el de fp32 salvo por el último paso de redondeo. fp16 guarda 10 bits y solo cinco de exponente, así que no pasa de 65 504 y los gradientes pequeños se quedan en cero por underflow si antes no se escala la pérdida. De ahí que el entrenamiento se haya pasado a bf16 y que para servir el modelo se recurra a int8. Cada formato es una decisión sobre qué precisión no le hace falta a un tensor.
Al servir el modelo, el eje que controlas es el del batch, y mueve dos números en sentidos opuestos. He cronometrado en la CPU de mi portátil un matmul de [B, 4096] por [4096, 4096], que cuesta 447 µs por fila con B = 1 y 22 µs por fila con B = 512, es decir, veinte veces más throughput. La llamada en sí ha pasado de 0,45 ms a 11,4 ms, y cada petición ha esperado veinticinco veces más, sin contar lo que haya tardado en llenarse el batch. Con cinco repeticiones en un solo portátil, lo que vale es la proporción y los dígitos exactos importan poco. El batching dinámico es el ajuste que se mueve entre esos dos números, y dónde lo pongas es una decisión de producto disfrazada de forma de tensor.
Qué vigila el ingeniero de ML
Son cuatro hábitos, aunque los dos primeros son en realidad uno solo.
Anota las formas en los límites de las funciones. jaxtyping con beartype convierte la firma en una comprobación que se ejecuta en cada llamada:
@jaxtyped(typechecker=beartype)def pool(x: Float[Tensor, "B T d"], mask: Bool[Tensor, "B T"]) -> Float[Tensor, "B d"]: m = mask[:, :, None].float() return (x * m).sum(1) / m.sum(1)Si dentro de la función reduces por el eje equivocado, salta un error al devolver el resultado, porque [T, d] no es [B, d]. Y si pasas al error cuadrático el objetivo con forma de columna, la comprobación falla antes de que la regla pueda estirar nada. Así se detectan en el límite dos de los bugs de arriba (la reducción, solo cuando B ≠ T), con la forma escrita justo donde va a mirar el siguiente que lea el código. Para el bug del batch de dieciséis, y para la reducción por el eje equivocado cuando B es igual a T, hace falta el hábito siguiente.
Prueba con tamaños que no puedan coincidir. La anotación no ve el bug del batch de dieciséis. Cuando B y T valen 16, los dos nombres se ligan al mismo número y la comprobación pasa. Con la reducción por el eje equivocado ocurre lo mismo, porque [T, d] y [B, d] son el mismo [16, 64]. El arreglo va en los tests, con B = 3, T = 5 y d = 7, y el código se queda como está. Cada None olvidado y cada confusión entre B, T y d dan entonces un tamaño que no coincide con nada, y el runtime lanza un error. Los ejes de tamaño 1 se siguen estirando con cualquier tamaño y las dos T de las puntuaciones son siempre iguales; esos casos necesitan la anotación o un test que compruebe valores. Los números redondos en los tests son los que dejan que el bug del batch de dieciséis llegue a producción. Por eso los dos primeros hábitos son uno solo: la anotación nunca es más estricta que los tamaños que ve.
La máscara y el lado del padding son una sola decisión. Apunta el lado junto a la configuración del tokenizador, calcula el índice del último token a partir de la máscara (nunca de T) y rechaza cualquier batch cuya máscara no cuadre con su padding. El código de pooling de la parte 2 ya divide por la suma de la máscara; asegúrate de que el tuyo también lo haga.
Calcula lo que cuesta un tensor antes de reservarlo. Multiplica la forma por los bytes de un número, para el tamaño de batch y el formato con los que vayas a servir el modelo. La demo de arriba hace esa multiplicación, y el resultado lo conoces antes de cargar un solo peso.
Esto tampoco es trabajo de modelado. La forma es el único tipo que este stack no comprueba nunca. Si los dispositivos no coinciden, salta un error, pero un eje equivocado se resuelve por broadcasting, y esa comprobación te toca escribirla a ti.