Skills › Scientific Agent Skills › pytorch-lightning › quick_trainer_setup.py
quick_trainer_setup.py
483 lines · skills/pytorch-lightning/scripts/quick_trainer_setup.py at 5f2bc03
1"""
2Quick Trainer Setup Examples for PyTorch Lightning.
4Targets Lightning 2.6.6. Helpers use local loggers and monitor val/loss, matching
5template_lightning_module.py. GPU/FSDP/DeepSpeed configurations are illustrative;
6CPU tests exercise the single-device helper with explicit CPU/full-precision overrides.
7This script provides Trainer configurations for common use cases.
8Copy and modify these configurations for your specific needs.
9"""
11import lightning as L
12from lightning.pytorch.callbacks import (
13 ModelCheckpoint,
14 EarlyStopping,
15 LearningRateMonitor,
16)
17from lightning.pytorch import loggers as pl_loggers
18from lightning.pytorch.strategies import DDPStrategy, FSDPStrategy
21# =============================================================================
22# 1. BASIC TRAINING (Single GPU/CPU)
23# =============================================================================
25def basic_trainer():
26 """
27 Simple trainer for quick prototyping.
28 Use for: Small models, debugging, single GPU training
29 """
30 trainer = L.Trainer(
31 max_epochs=10,
32 accelerator="auto", # Automatically select GPU/CPU
33 devices=1, # Keep the basic smoke run on one device
34 enable_progress_bar=True,
35 logger=pl_loggers.CSVLogger("logs", name="basic"),
36 )
37 return trainer
40# =============================================================================
41# 2. DEBUGGING CONFIGURATION
42# =============================================================================
44def debug_trainer():
45 """
46 Trainer for debugging with fast dev run and anomaly detection.
47 Use for: Finding bugs, testing code quickly
48 """
49 trainer = L.Trainer(
50 fast_dev_run=True, # Run 1 batch through train/val/test
51 accelerator="cpu", # Use CPU for easier debugging
52 detect_anomaly=True, # Detect NaN/Inf in gradients
53 log_every_n_steps=1, # Log every step
54 enable_progress_bar=True,
55 )
56 return trainer
59# =============================================================================
60# 3. PRODUCTION TRAINING (Single GPU)
61# =============================================================================
63def production_single_gpu_trainer(
64 max_epochs=100,
65 log_dir="logs",
66 checkpoint_dir="checkpoints",
67 accelerator="gpu",
68 precision="16-mixed",
69):
70 """
71 Production-ready trainer for single GPU with checkpointing and logging.
72 Use for: Final training runs on single GPU
73 """
74 # Callbacks
75 checkpoint_callback = ModelCheckpoint(
76 dirpath=checkpoint_dir,
77 filename="epoch-{epoch:02d}-step-{step:06d}",
78 auto_insert_metric_name=False,
79 monitor="val/loss",
80 mode="min",
81 save_top_k=3,
82 save_last=True,
83 verbose=True,
84 )
86 early_stop_callback = EarlyStopping(
87 monitor="val/loss",
88 patience=10,
89 mode="min",
90 verbose=True,
91 )
93 lr_monitor = LearningRateMonitor(logging_interval="step")
95 # Logger
96 csv_logger = pl_loggers.CSVLogger(
97 save_dir=log_dir,
98 name="my_model",
99 )
101 # Trainer
102 trainer = L.Trainer(
103 max_epochs=max_epochs,
104 accelerator=accelerator,
105 devices=1,
106 precision=precision, # Mixed precision for speed
107 callbacks=[
108 checkpoint_callback,
109 early_stop_callback,
110 lr_monitor,
111 ],
112 logger=csv_logger,
113 log_every_n_steps=50,
114 gradient_clip_val=1.0, # Clip gradients
115 enable_progress_bar=True,
116 )
118 return trainer
121# =============================================================================
122# 4. MULTI-GPU TRAINING (DDP)
123# =============================================================================
125def multi_gpu_ddp_trainer(
126 max_epochs=100,
127 num_gpus=4,
128 log_dir="logs",
129 checkpoint_dir="checkpoints"
130):
131 """
132 Multi-GPU training with Distributed Data Parallel.
133 Use when a full replica plus optimizer state and activations fits per GPU
134 """
135 # Callbacks
136 checkpoint_callback = ModelCheckpoint(
137 dirpath=checkpoint_dir,
138 filename="epoch-{epoch:02d}-step-{step:06d}",
139 auto_insert_metric_name=False,
140 monitor="val/loss",
141 mode="min",
142 save_top_k=3,
143 save_last=True,
144 )
146 early_stop_callback = EarlyStopping(
147 monitor="val/loss",
148 patience=10,
149 mode="min",
150 )
152 lr_monitor = LearningRateMonitor(logging_interval="step")
154 # Logger
155 csv_logger = pl_loggers.CSVLogger(
156 name="my-project",
157 save_dir=log_dir,
158 )
160 # Trainer
161 trainer = L.Trainer(
162 max_epochs=max_epochs,
163 accelerator="gpu",
164 devices=num_gpus,
165 strategy=DDPStrategy(
166 find_unused_parameters=False,
167 gradient_as_bucket_view=True,
168 ),
169 precision="16-mixed",
170 callbacks=[
171 checkpoint_callback,
172 early_stop_callback,
173 lr_monitor,
174 ],
175 logger=csv_logger,
176 log_every_n_steps=50,
177 gradient_clip_val=1.0,
178 sync_batchnorm=True, # Sync batch norm across GPUs
179 )
181 return trainer
184# =============================================================================
185# 5. LARGE MODEL TRAINING (FSDP)
186# =============================================================================
188def large_model_fsdp_trainer(
189 max_epochs=100,
190 num_gpus=8,
191 log_dir="logs",
192 checkpoint_dir="checkpoints"
193):
194 """
195 Training for memory-constrained models with FSDP.
196 Use for: Large transformers, models that don't fit in single GPU
197 """
198 import torch.nn as nn
200 # Callbacks
201 checkpoint_callback = ModelCheckpoint(
202 dirpath=checkpoint_dir,
203 filename="epoch-{epoch:02d}-step-{step:06d}",
204 auto_insert_metric_name=False,
205 monitor="val/loss",
206 mode="min",
207 save_top_k=3,
208 save_last=True,
209 )
211 lr_monitor = LearningRateMonitor(logging_interval="step")
213 # Logger
214 csv_logger = pl_loggers.CSVLogger(
215 name="large-model",
216 save_dir=log_dir,
217 )
219 # Trainer with FSDP
220 trainer = L.Trainer(
221 max_epochs=max_epochs,
222 accelerator="gpu",
223 devices=num_gpus,
224 strategy=FSDPStrategy(
225 sharding_strategy="FULL_SHARD",
226 activation_checkpointing_policy={
227 nn.TransformerEncoderLayer,
228 nn.TransformerDecoderLayer,
229 },
230 cpu_offload=False, # Set True if GPU memory insufficient
231 ),
232 precision="bf16-mixed", # BFloat16 for A100/H100
233 callbacks=[
234 checkpoint_callback,
235 lr_monitor,
236 ],
237 logger=csv_logger,
238 log_every_n_steps=10,
239 gradient_clip_val=1.0,
240 accumulate_grad_batches=4, # Gradient accumulation
241 )
243 return trainer
246# =============================================================================
247# 6. VERY LARGE MODEL TRAINING (DeepSpeed)
248# =============================================================================
250def deepspeed_trainer(
251 max_epochs=100,
252 num_gpus=8,
253 stage=3,
254 log_dir="logs",
255 checkpoint_dir="checkpoints"
256):
257 """
258 Training for very large models with DeepSpeed.
259 Use when ZeRO sharding is needed; measure memory and throughput
260 """
261 if stage not in (1, 2, 3):
262 raise ValueError("DeepSpeed stage must be 1, 2, or 3")
264 # Callbacks
265 checkpoint_callback = ModelCheckpoint(
266 dirpath=checkpoint_dir,
267 filename="{epoch:02d}-{step:06d}",
268 monitor="val/loss", # Required: save_top_k > 1 needs a quantity
269 mode="min",
270 save_top_k=3,
271 save_last=True,
272 every_n_epochs=1, # Rank using fresh validation metrics
273 )
275 lr_monitor = LearningRateMonitor(logging_interval="step")
277 # Logger
278 csv_logger = pl_loggers.CSVLogger(
279 name="very-large-model",
280 save_dir=log_dir,
281 )
283 # Select DeepSpeed stage
284 strategy = f"deepspeed_stage_{stage}"
286 # Trainer
287 trainer = L.Trainer(
288 max_epochs=max_epochs,
289 accelerator="gpu",
290 devices=num_gpus,
291 strategy=strategy,
292 precision="16-mixed",
293 callbacks=[
294 checkpoint_callback,
295 lr_monitor,
296 ],
297 logger=csv_logger,
298 log_every_n_steps=10,
299 gradient_clip_val=1.0,
300 accumulate_grad_batches=4,
301 )
303 return trainer
306# =============================================================================
307# 7. HYPERPARAMETER TUNING
308# =============================================================================
310def hyperparameter_tuning_trainer(max_epochs=50):
311 """
312 Lightweight trainer for hyperparameter search.
313 Use for: Quick experiments, hyperparameter tuning
314 """
315 trainer = L.Trainer(
316 max_epochs=max_epochs,
317 accelerator="auto",
318 devices=1,
319 enable_checkpointing=False, # Don't save checkpoints
320 logger=False, # Disable logging
321 enable_progress_bar=False,
322 limit_train_batches=0.5, # Use 50% of training data
323 limit_val_batches=0.5, # Use 50% of validation data
324 )
325 return trainer
328# =============================================================================
329# 8. OVERFITTING TEST
330# =============================================================================
332def overfit_test_trainer(num_batches=10):
333 """
334 Trainer for overfitting on small subset to verify model capacity.
335 Use for: Testing if model can learn, debugging
336 """
337 trainer = L.Trainer(
338 max_epochs=100,
339 accelerator="auto",
340 devices=1,
341 overfit_batches=num_batches, # Overfit on N batches
342 log_every_n_steps=1,
343 enable_progress_bar=True,
344 )
345 return trainer
348# =============================================================================
349# 9. TIME-LIMITED TRAINING (SLURM)
350# =============================================================================
352def time_limited_trainer(
353 max_time_hours=23.5,
354 max_epochs=1000,
355 checkpoint_dir="checkpoints"
356):
357 """
358 Training with time limit for SLURM clusters.
359 Use for: Cluster jobs with time limits
360 """
361 from datetime import timedelta
363 checkpoint_callback = ModelCheckpoint(
364 dirpath=checkpoint_dir,
365 monitor="val/loss", # Required: save_top_k > 1 needs a quantity
366 mode="min",
367 save_top_k=3,
368 save_last=True, # Important for resuming
369 every_n_epochs=1,
370 )
372 trainer = L.Trainer(
373 max_epochs=max_epochs,
374 max_time=timedelta(hours=max_time_hours),
375 accelerator="gpu",
376 devices="auto",
377 callbacks=[checkpoint_callback],
378 log_every_n_steps=50,
379 )
381 return trainer
384# =============================================================================
385# 10. REPRODUCIBLE RESEARCH
386# =============================================================================
388def reproducible_trainer(seed=42, max_epochs=100):
389 """
390 Seeded trainer for repeatable runs on the same software/hardware stack.
391 Use for: Publications, reproducible results
392 """
393 # Call this before creating the model and DataModule.
394 L.seed_everything(seed, workers=True)
396 # Callbacks
397 checkpoint_callback = ModelCheckpoint(
398 dirpath="checkpoints",
399 filename="epoch-{epoch:02d}-step-{step:06d}",
400 auto_insert_metric_name=False,
401 monitor="val/loss",
402 mode="min",
403 save_top_k=3,
404 save_last=True,
405 )
407 # Trainer
408 trainer = L.Trainer(
409 max_epochs=max_epochs,
410 accelerator="gpu",
411 devices=1,
412 precision="32-true", # Full precision for reproducibility
413 deterministic=True, # Use deterministic algorithms
414 benchmark=False, # Disable cudnn benchmarking
415 callbacks=[checkpoint_callback],
416 log_every_n_steps=50,
417 )
419 return trainer
422# =============================================================================
423# USAGE EXAMPLES
424# =============================================================================
426if __name__ == "__main__":
427 from lightning.pytorch.accelerators import CUDAAccelerator
429 # Trainer construction validates the requested hardware, so the GPU-only
430 # configurations below can only be built where that hardware exists.
431 cuda_devices = CUDAAccelerator.auto_device_count() if CUDAAccelerator.is_available() else 0
433 print("PyTorch Lightning Trainer Configurations\n")
435 # Example 1: Basic training
436 print("1. Basic Trainer:")
437 trainer = basic_trainer()
438 print(f" - Max epochs: {trainer.max_epochs}")
439 print(f" - Accelerator: {trainer.accelerator}")
440 print()
442 # Example 2: Debug training
443 print("2. Debug Trainer:")
444 trainer = debug_trainer()
445 print(f" - Fast dev run: {trainer.fast_dev_run}")
446 print(f" - Accelerator: {trainer.accelerator}")
447 print()
449 # Example 3: Production single GPU
450 print("3. Production Single GPU Trainer:")
451 if cuda_devices >= 1:
452 trainer = production_single_gpu_trainer(max_epochs=100)
453 print(f" - Max epochs: {trainer.max_epochs}")
454 print(f" - Precision: {trainer.precision}")
455 print(f" - Callbacks: {len(trainer.callbacks)}")
456 else:
457 print(" - skipped: no CUDA GPU on this machine")
458 print()
460 # Example 4: Multi-GPU DDP
461 print("4. Multi-GPU DDP Trainer:")
462 if cuda_devices >= 4:
463 trainer = multi_gpu_ddp_trainer(num_gpus=4)
464 print(f" - Strategy: {trainer.strategy}")
465 print(f" - Devices: {trainer.num_devices}")
466 else:
467 print(f" - skipped: needs 4 CUDA GPUs, found {cuda_devices}")
468 print()
470 # Example 5: FSDP for large models
471 print("5. FSDP Trainer for Large Models:")
472 if cuda_devices >= 8:
473 trainer = large_model_fsdp_trainer(num_gpus=8)
474 print(f" - Strategy: {trainer.strategy}")
475 print(f" - Precision: {trainer.precision}")
476 else:
477 print(f" - skipped: needs 8 CUDA GPUs, found {cuda_devices}")
478 print()
480 print("\nTo use these configurations:")
481 print("1. Import the desired function")
482 print("2. Create trainer: trainer = production_single_gpu_trainer()")
483 print("3. Train model: trainer.fit(model, datamodule=dm)")