1 lightning.training
lightning.training
Training machinery shared by one-step VIN and finite-horizon QH.
1.1 Package guide
This package owns trainer construction, callbacks, optimizer/scheduler configuration and gradient observation shared by one_step and finite_horizon. Each task retains data admission, loss/support semantics, checkpoint selection and experiment transaction ordering.
1.2 Construction and lifecycle
TrainerFactoryConfig constructs a Lightning trainer after resolving logger and callback ownership. It accepts an explicit Optuna configuration/trial; the existing experiment-config fallback remains available for one-step callers. Debug mode uses Lightning’s default logger; configured normal runs use W&B or no logger. Checkpointing follows the actual callback set. Float32 matmul precision is process-wide when configured.
TrainerCallbacksConfig constructs checkpoint, progress, learning rate, pruning and optional optimizer-diagnostic callbacks. optimizers.py owns typed AdamW, OneCycle and ReduceLROnPlateau configuration/factories. The task module supplies optimizer parameters and monitor/update semantics. Experiments construct trainer, data-module and module owners; Lightning invokes task hooks to construct optimizers and schedulers through these shared factories. Lightning owns runtime lifecycle and teardown.
1.3 Research execution policy
execution_profile.py owns the shared CUDA-AMP and MPS-FP32 profiles used by current Q1/QH launchers and the package-owned Mac campaign. It applies explicit trainer/optimizer controls, admits hardware and precision, observes named CPU fallback, and binds immutable W&B attempt segments. It does not launch a fit or create a W&B run. Attempt writer locks and publication readback retain their existing local and cloud semantics.
New freeze protocols declare version 2 and an execution_policy_source record naming the package source path, record version and exact file digest. Historical Mac campaign version 1 keeps execution_profile_sha256 bound to the original script at the explicitly named producing checkout; observers authenticate those original bytes rather than reinterpret the field as the moved package owner. Existing frozen protocols and measured publications are unchanged.
1.4 Optimizer and gradient observation
Gradient-norm logging lives in gradient_logging.py because it observes Lightning’s optimizer lifecycle. It reads detached gradients and does not mutate autograd state. The configured depth/include/exclude rules bound the number of logged parameter groups. The same owner encodes finite, detached, regularly spaced histogram samples with a bounded CPU payload. Callers keep logger selection and cadence: QH observes successful-update counts, while the shared callback uses Lightning’s attempted-step count. The helper does not query scorers, retain activation graphs, or decide whether AMP stepped.
For manual fitted-Q runs, opt into trainer.callbacks.optimizer_diagnostics.level = "summary" (or "detailed"). The shared callback observes unscaled gradients at on_before_optimizer_step; summary logs the global L2 norm, parameters with/without gradients, and capped parameter-group learning rates and weight decay. Detailed adds at most max_items parameter norms. every_n_steps controls observation cadence; "off" attaches no callback. Only optimizer-owned parameters enter diagnostics, so the fitted-Q target network is excluded. Values are rank-local scalar observations, not proof that GradScaler completed an update. No tensors or graphs are retained. Automatic clipping follows this hook. QH norm clipping runs in the module’s corresponding hook, after callback observation and AMP unscaling. Other manual modules determine their own clipping order. Under AMP, Lightning’s global_step counts optimizer attempts even when GradScaler skips an overflow. QH logs train/optimizer_update_skipped and advances its scheduler, target-copy cadence and trained-support counts only after a successful raw update. Callback cadence follows Lightning attempts; QH output diagnostics follow the module’s completed-update counter. Each optimizer callback payload carries train-optimizer/observed_attempted_step: the zero-based trainer.global_step at the pre-step observation. Carried metrics keep that original coordinate; later transport rows are not new observations. The coordinate counts attempted steps, including AMP-skipped attempts, rather than completed raw updates.
1.5 Verification
Use trainer factory, checkpoint callbacks, optimizer validation, plateau scheduling, gradient logging, optimizer diagnostics and W&B logging for the affected seam. Task lifecycle tests verify how each experiment composes this machinery. Detailed configuration defaults and callable contracts remain in source.