A lightweight wrapper that routes a linear layer through TransformerEngine.
| 188 | |
| 189 | |
| 190 | class TransformerEngineLinear(nn.Module): |
| 191 | """A lightweight wrapper that routes a linear layer through TransformerEngine.""" |
| 192 | |
| 193 | def __init__( |
| 194 | self, |
| 195 | module: nn.Linear, |
| 196 | module_name: str, |
| 197 | module_config: Any, |
| 198 | inference_only: bool = False, |
| 199 | low_precision_weights: bool = False, |
| 200 | te_recipe_kwargs: dict[str, Any] | None = None, |
| 201 | te_module_kwargs: dict[str, Any] | None = None, |
| 202 | ) -> None: |
| 203 | super().__init__() |
| 204 | |
| 205 | try: |
| 206 | te = importlib.import_module("transformer_engine.pytorch") |
| 207 | except ImportError as exc: |
| 208 | raise ImportError( |
| 209 | "TransformerEngine is not installed, but `use_transformer_engine=True` " |
| 210 | "was requested." |
| 211 | ) from exc |
| 212 | |
| 213 | if module.weight.device.type != "cuda": |
| 214 | raise ValueError( |
| 215 | "TransformerEngine replacement requires CUDA modules. " |
| 216 | f"Module `{module_name}` is on `{module.weight.device}`." |
| 217 | ) |
| 218 | |
| 219 | self.module_name = module_name |
| 220 | self.in_features = module.in_features |
| 221 | self.out_features = module.out_features |
| 222 | self.inference_only = inference_only |
| 223 | self.low_precision_weights = low_precision_weights |
| 224 | self._te = te |
| 225 | self._recipe = _build_te_recipe( |
| 226 | module_config=module_config, |
| 227 | te_recipe_kwargs=te_recipe_kwargs, |
| 228 | ) |
| 229 | |
| 230 | module_kwargs = dict(te_module_kwargs or {}) |
| 231 | module_kwargs.setdefault("device", module.weight.device) |
| 232 | module_kwargs.setdefault("params_dtype", module.weight.dtype) |
| 233 | module_kwargs.setdefault("name", module_name) |
| 234 | |
| 235 | fp8_model_init_fn = getattr(te, "fp8_model_init", None) |
| 236 | if self.low_precision_weights and fp8_model_init_fn is None: |
| 237 | warnings.warn( |
| 238 | "TransformerEngine low-precision parameter init requested, but " |
| 239 | "`fp8_model_init` is unavailable. Falling back to regular TE parameter " |
| 240 | "storage for this inference path.", |
| 241 | stacklevel=2, |
| 242 | ) |
| 243 | self.low_precision_weights = False |
| 244 | |
| 245 | model_init_context = ( |
| 246 | fp8_model_init_fn( |
| 247 | enabled=True, |
no outgoing calls
no test coverage detected