Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / trainer / config.py: 98%
305 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-13 05:07 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""Training configuration schema — strict three-tier (model/data/train).
17Top-level keys are exactly ``model``, ``data`` and ``train`` and nothing else;
18anything outside this three-tier schema is rejected by the parser.
20YAML shape::
22 model:
23 name: qwen3_5
24 weights_path: /path/to/weights
25 data:
26 type: hf_datasets
27 train_path: /path/to/data
28 train:
29 max_steps: 100
30 micro_batch_size: 1
31 global_batch_size: 8
32 seed: 42
33 init_device: meta
34 optimizer:
35 type: adamw
36 lr: 1.0e-4
37 accelerator:
38 dp_shard: 8
39 tp: 1
40 mixed_precision:
41 enabled: true
42 param_dtype: bfloat16
43 checkpoint:
44 output_dir: outputs/run1
45 ...
46"""
47import argparse
48import difflib
49import logging
50import os
51from dataclasses import dataclass, field, fields, is_dataclass
52from typing import Any, Dict, List, Optional, Type, TypeVar, Union, get_args, get_origin
54import yaml
56logger = logging.getLogger(__name__)
58T = TypeVar("T")
60_BOOL_TRUE_STRINGS = frozenset(("true", "yes", "y", "on", "1", "t"))
61_BOOL_FALSE_STRINGS = frozenset(("false", "no", "n", "off", "0", "f"))
63# ============================================================================
64# model:
65# ============================================================================
68@dataclass
69class ModelConfig:
70 """``model.*`` — model identity, weights, and architecture overrides.
72 Only universal transformer fields are typed here. Anything model-
73 specific (mRoPE section split, MoE expert geometry, linear-attention
74 head counts, ``layer_types`` ...) goes through ``config_overrides`` —
75 a free-form ``dict`` that is merged into the underlying model
76 constructor by the model's ``build_model_fn``.
77 """
78 name: str = "qwen3_5"
79 weights_path: Optional[str] = None
80 tokenizer_path: Optional[str] = None
81 freeze_modules: Optional[list] = None
82 tp_plan: Optional[dict] = None
83 cp_modules: Optional[list] = None
84 ep_modules: Optional[list] = None
85 # Universal transformer architecture overrides.
86 num_hidden_layers: Optional[int] = None
87 hidden_size: Optional[int] = None
88 intermediate_size: Optional[int] = None
89 num_attention_heads: Optional[int] = None
90 num_key_value_heads: Optional[int] = None
91 vocab_size: Optional[int] = None
92 max_position_embeddings: Optional[int] = None
93 # Free-form per-model overrides handed to ``build_model_fn``.
94 config_overrides: Optional[dict] = None
96# ============================================================================
97# data:
98# ============================================================================
101@dataclass
102class DataConfig:
103 """``data.*`` — dataset, tokenizer/processor, sampler, batch shape.
105 - ``streaming``: only ``False`` supported today.
106 - ``num_workers``: keep ≥ 2 for real datasets.
107 - ``shuffle``: when ``False``, sampler reads samples in dataset order.
108 """
109 type: str = "dummy"
110 train_path: Optional[Any] = None
111 subset: Optional[str] = None
112 max_seq_len: int = 2048
113 text_key: str = "text"
114 train_size: Optional[int] = None
115 # multimodal / VL (synthetic vl_dummy path)
116 template: str = "empty"
117 image_key: str = "image"
118 messages_key: str = "messages"
119 image_token_id: int = 151655
120 video_token_id: int = 151656
121 vl_video: bool = False
122 vl_grid_t: int = 2
123 vl_grid_h: int = 2
124 vl_grid_w: int = 2
125 # loader perf
126 streaming: bool = False
127 num_workers: int = 0
128 prefetch_factor: Optional[int] = None
129 pin_memory: bool = True
130 shuffle: bool = True
131 # Megatron .bin/.idx options (data.type='megatron'). ``train_path`` is a
132 # single path-prefix, or a blend spec ("w1 prefix1 w2 prefix2 ...") /
133 # list-of-pairs. ``megatron_seed`` defaults to ``train.seed`` when None.
134 megatron_seed: Optional[int] = None
135 pad_token_id: int = 0
136 eod_token_id: Optional[int] = None
137 eod_mask_loss: bool = False
138 mmap_bin_files: bool = True
140# ============================================================================
141# train.* — sub-configs
142# ============================================================================
145@dataclass
146class AcceleratorConfig:
147 """``train.accelerator.*`` — parallelism topology.
149 Two ways to express data parallelism:
151 - **Legacy single field** (back-compat): ``dp`` only — maps to
152 ``dp_shard`` for FSDP.
153 - ** split**: ``dp_replicate`` (HSDP outer) and
154 ``dp_shard`` (FSDP inner). Pass ``dp_shard=-1`` to auto-fill from
155 ``world_size / (dp_replicate * cp * tp * pp)``.
157 For MoE: ``etp`` controls expert TP. Must equal ``tp`` or ``1``.
158 ``moe_token_dispatcher_type`` selects the EP token exchange algorithm.
159 ``npu_nums_per_device`` is the inner-EP degree for the deredundency
160 dispatcher; ``oep`` is inferred as ``ep // npu_nums_per_device``.
161 """
162 dp: Optional[int] = None
163 dp_replicate: int = 1
164 dp_shard: Optional[int] = None
165 tp: int = 1
166 cp: int = 1
167 pp: int = 1
168 # Number of pipeline micro-batches per optimizer step when ``pp > 1``.
169 # The global batch is split into this many micro-batches along dim 0 and
170 # streamed through the pipeline stages; ``global_batch_size`` must be
171 # divisible by it. Defaults to 1 (no micro-batching).
172 pp_micro_batch_num: int = 1
173 # Pipeline schedule when ``pp > 1``: ``"gpipe"`` (all-forward then
174 # all-backward) or ``"1f1b"`` (one-forward-one-backward steady state).
175 # ``None`` keeps each model's default (dense → gpipe, MoE → 1f1b).
176 pp_schedule: Optional[str] = None
177 # Virtual-pipeline (VPP) degree: number of non-contiguous stage chunks each
178 # PP rank owns. ``1`` (default) is the plain single-stage-per-rank pipeline.
179 # ``>1`` builds ``pp * pp_vpp`` interleaved global stages (rank ``r`` owns
180 # stages ``r, r+pp, r+2*pp, ...``) and drives them with interleaved 1F1B.
181 pp_vpp: int = 1
182 # Per-global-stage decoder-layer counts (length ``pp * pp_vpp``, summing to
183 # ``num_hidden_layers``). ``None`` (default) keeps the even split with the
184 # remainder on the later stages. Lets users rebalance stages whose extra
185 # modules (embed / visual tower / lm_head) dominate memory or latency.
186 pp_layer_split: Optional[List[int]] = None
187 ep: int = 1
188 etp: int = 1
189 moe_token_dispatcher_type: str = "all_to_all"
190 npu_nums_per_device: int = 8
191 zero_stage: int = 0
192 reshard_after_forward: bool = True
193 async_cp: bool = False
194 ulysses_degree: Optional[int] = None
195 # Bucketed reduce-scatter: single fused RS per FSDP unit, stable fp32
196 # reduction order across runs.
197 comm_fusion: bool = True
198 # Offload sharded params, grads, and optimizer states to CPU so large
199 # checkpoints can fit without device-resident master weights and Adam state.
200 cpu_offload: bool = False
203@dataclass
204class MixedPrecisionConfig:
205 """``train.mixed_precision.*`` — mixed-precision forward configuration.
207 FSDP2 ``MixedPrecisionPolicy``: ``param_dtype`` is the all-gather'd
208 forward dtype, ``reduce_dtype`` is the reduce-scatter dtype for grads,
209 and ``output_dtype`` controls the forward-output dtype at FSDP wrap
210 boundaries (leave ``None`` to inherit from ``param_dtype``).
211 """
212 enabled: bool = False
213 param_dtype: str = "bfloat16"
214 reduce_dtype: str = "float32"
215 output_dtype: Optional[str] = None
218@dataclass
219class GradientCheckpointingConfig:
220 """``train.gradient_checkpointing.*`` — activation recomputation.
222 . Modes: ``"off"``, ``"full"``, or ``"selective"``.
223 """
224 activation_checkpoint: str = "off"
227@dataclass
228class OptimizerConfig:
229 """``train.optimizer.*`` — optimizer + LR schedule + grad clip.
231 ``loss_aggregation``: how the per-micro-batch loss is scaled before
232 backward. ``"token_weighted"`` divides the summed loss by the global
233 valid-token count; ``"rank_average"`` averages per-rank micro-batch
234 means and is preferred when batches have variable valid-token counts
235 across ranks.
236 ``max_grad_norm``: values <= 0 disable gradient clipping.
237 """
238 type: str = "adamw"
239 lr: float = 1e-4
240 lr_min: float = 1e-5
241 lr_decay_style: str = "cosine"
242 lr_warmup_ratio: float = 0.1
243 loss_aggregation: str = "token_weighted"
244 weight_decay: float = 0.01
245 max_grad_norm: float = 1.0
246 bsz_warmup_ratio: float = 0.0
247 eps: float = 1e-8
248 betas: tuple = (0.9, 0.999)
249 # ``None`` lets torch pick the foreach kernel; set ``False`` in YAML when a
250 # run must force the deterministic per-parameter loop.
251 foreach: Optional[bool] = None
254@dataclass
255class CheckpointConfig:
256 """``train.checkpoint.*`` — DCP save / load + HF export."""
257 output_dir: str = "outputs"
258 save_steps: int = 500
259 save_hf_weights: bool = True
260 load_path: Optional[str] = None
261 save_async: bool = False
264@dataclass
265class LoggingConfig:
266 """``train.logging.*`` — console / metric output (consumed by LoggingCallback)."""
267 report_to: str = "none"
268 report_global_loss: bool = False
269 log_steps: int = 10
270 report_throughput: bool = True
271 model_flops_per_token: Optional[int] = None
272 peak_tflops: Optional[float] = None # e.g. 312.0 for A100 bf16
275@dataclass
276class TensorBoardConfig:
277 """``train.tensorboard.*`` — TB SummaryWriter on rank 0."""
278 enabled: bool = False
279 output_dir: str = "tb_traces"
280 log_steps: int = 1
283@dataclass
284class WandbConfig:
285 """``train.wandb.*`` — W&B run logging on rank 0."""
286 enabled: bool = False
287 project: str = "hyper-parallel"
288 run_name: Optional[str] = None
289 log_steps: int = 1
292@dataclass
293class ProfileConfig:
294 """``train.profile.*`` — torch.profiler schedule ().
296 Schedule semantics: wait → warmup → active.
297 """
298 enabled: bool = False
299 output_dir: str = "profiler_traces"
300 wait_steps: int = 1
301 warmup_steps: int = 1
302 active_steps: int = 3
305@dataclass
306class MemoryMonitorConfig:
307 """``train.memory_monitor.*`` — periodic device-memory snapshot."""
308 enabled: bool = False
309 log_steps: int = 50
310 reset_peak_each_step: bool = False
313@dataclass
314class MoEMonitorConfig:
315 """``train.moe_monitor.*`` — MoE routing / load-balance monitor.
317 When enabled, :class:`~hyper_parallel.core.moe_utils.MoEMonitorCallback`
318 automatically syncs ``tokens_per_expert`` across distributed ranks and
319 updates ``expert_bias`` after each optimizer step. The mean ``aux_loss``
320 across MoE layers is exposed via ``last_mean_aux_loss`` so that
321 :class:`LoggingCallback` can print it alongside the main training loss.
323 DP/TP+SP/CP group information is automatically obtained from the trainer's
324 device mesh — no manual group configuration needed.
326 Args:
327 enabled: Whether to activate the MoE monitor callback.
328 lr: Step size for expert bias updates. Defaults to ``1e-3``.
329 num_recomputations: Number of forward executions per optimizer step.
330 Default ``1``. Set to ``2`` when activation checkpoint is enabled.
331 """
332 enabled: bool = False
333 lr: float = 1e-3
334 num_recomputations: int = 1
337@dataclass
338class EvalConfig:
339 """``train.eval.*`` — eval cadence + dataset."""
340 eval_steps: int = 0
341 eval_dataset: Optional[str] = None
344@dataclass
345class DebugConfig:
346 """``train.debug.*`` — reproducibility and numerical-stability knobs.
348 All flags here are production-safe; they tune determinism (CI / paper
349 reproducibility), guard against numerical blow-ups, and bound memory
350 growth in long runs.
351 """
352 deterministic: bool = False
353 deterministic_warn_only: bool = False
354 check_nan_inf: bool = False
355 gc_steps: int = 0
357# ============================================================================
358# train: (top of the train section, holds the sub-configs)
359# ============================================================================
362@dataclass
363class TrainConfig:
364 """``train.*`` — full training-section config.
366 Flat fields cover the basic loop knobs (steps, batch shape, init device,
367 seed, comm backend); nested sub-configs cover everything else.
368 """
369 # Loop shape
370 max_steps: int = 100
371 num_train_epochs: int = 1
372 global_batch_size: int = 8
373 micro_batch_size: int = 1
374 seed: int = 42
376 # Runtime / device
377 backend: str = "torch"
378 init_device: str = "meta"
379 comm_backend: Optional[str] = None
380 local_rank: int = 0 # set from LOCAL_RANK env by parser
382 # Sub-configs
383 accelerator: AcceleratorConfig = field(default_factory=AcceleratorConfig)
384 mixed_precision: MixedPrecisionConfig = field(default_factory=MixedPrecisionConfig)
385 gradient_checkpointing: GradientCheckpointingConfig = field(
386 default_factory=GradientCheckpointingConfig
387 )
388 optimizer: OptimizerConfig = field(default_factory=OptimizerConfig)
389 checkpoint: CheckpointConfig = field(default_factory=CheckpointConfig)
390 logging: LoggingConfig = field(default_factory=LoggingConfig)
391 tensorboard: TensorBoardConfig = field(default_factory=TensorBoardConfig)
392 wandb: WandbConfig = field(default_factory=WandbConfig)
393 profile: ProfileConfig = field(default_factory=ProfileConfig)
394 memory_monitor: MemoryMonitorConfig = field(default_factory=MemoryMonitorConfig)
395 moe_monitor: MoEMonitorConfig = field(default_factory=MoEMonitorConfig)
396 eval: EvalConfig = field(default_factory=EvalConfig)
397 debug: DebugConfig = field(default_factory=DebugConfig)
399# ============================================================================
400# Top-level: model / data / train (and only these three)
401# ============================================================================
404@dataclass
405class HyperTrainerConfig:
406 """Top-level config — strict three-tier ().
408 Allowed top-level keys: ``model``, ``data``, ``train``. Anything else in
409 the YAML is rejected by the parser with a typo-suggestion message.
410 """
411 model: ModelConfig = field(default_factory=ModelConfig)
412 data: DataConfig = field(default_factory=DataConfig)
413 train: TrainConfig = field(default_factory=TrainConfig)
415 # Computed (no user input)
416 train_steps: int = 0
418 def __post_init__(self):
419 self.train_steps = self.train.max_steps
421# ==============================================================================
422# CLI / YAML parser
423# ==============================================================================
424# Configuration parser: YAML file + CLI dot-path overrides.
425#
426# Supports:
427# - Unknown YAML/CLI keys emit a warning with difflib closest-match suggestions.
428# - Bool fields accept string aliases: ``true/yes/y/on/1/t`` -> ``True``,
429# ``false/no/n/off/0/f`` -> ``False``. Only applied when the dataclass
430# field type resolves to ``bool`` or ``Optional[bool]`` to avoid ambiguity.
433def _string_to_bool(value: Any) -> bool:
434 """Convert common string representations of booleans to ``bool``.
436 Accepts: ``true/yes/y/on/1/t`` → ``True``,
437 ``false/no/n/off/0/f`` → ``False``.
439 Args:
440 value: A string or bool value.
442 Returns:
443 The corresponding ``bool``.
445 Raises:
446 ValueError: When the string cannot be mapped to a bool.
447 """
448 if isinstance(value, bool):
449 return value
450 if isinstance(value, str):
451 lower = value.lower()
452 if lower in _BOOL_TRUE_STRINGS:
453 return True
454 if lower in _BOOL_FALSE_STRINGS:
455 return False
456 raise ValueError(
457 f"Cannot convert {value!r} to bool. "
458 "Expected one of: true/false/yes/no/y/n/on/off/1/0/t/f"
459 )
462def _resolve_field_type(cls: Type, dot_path: str) -> Optional[Type]:
463 """Walk a dataclass hierarchy to find the resolved type of a dot-path field.
465 Args:
466 cls: Root dataclass class.
467 dot_path: Dot-separated field path, e.g. ``"debug.deterministic"``.
469 Returns:
470 The resolved Python type, or ``None`` if the path cannot be resolved.
471 """
472 parts = dot_path.split(".")
473 current_cls = cls
474 for part in parts:
475 if not is_dataclass(current_cls):
476 return None
477 found = None
478 for f in fields(current_cls):
479 if f.name == part:
480 found = f
481 break
482 if found is None:
483 return None
484 field_type = found.type
485 # Unwrap Optional[X] → X
486 origin = get_origin(field_type)
487 if origin is Union:
488 unwrapped = [a for a in get_args(field_type) if a is not type(None)]
489 field_type = unwrapped[0] if len(unwrapped) == 1 else field_type
490 current_cls = field_type
491 return current_cls
494def _coerce_cli_value(raw: str, dot_path: str, root_class: Type) -> Any:
495 """Parse a CLI string value, coercing to the correct type for the field.
497 Bool fields accept an extended string set. For all other
498 fields the existing int → float → str heuristic is used.
500 Args:
501 raw: Raw string from the CLI.
502 dot_path: Dot-separated field path used for type lookup.
503 root_class: Root dataclass class for type resolution.
505 Returns:
506 Coerced value.
507 """
508 field_type = _resolve_field_type(root_class, dot_path)
509 if field_type is bool:
510 try:
511 return _string_to_bool(raw)
512 except ValueError:
513 pass # fall through to heuristic below
514 # Existing heuristic: int → float → bool-string → str
515 try:
516 return int(raw)
517 except ValueError:
518 pass
519 try:
520 return float(raw)
521 except ValueError:
522 pass
523 if raw.lower() in ("true", "false"):
524 return raw.lower() == "true"
525 return raw
528def _deep_update(source: Dict[str, Any], overrides: Dict[str, Any]) -> Dict[str, Any]:
529 """Recursively update source dict with overrides dict."""
530 for key, value in overrides.items():
531 if isinstance(value, dict) and isinstance(source.get(key), dict):
532 source[key] = _deep_update(source[key], value)
533 else:
534 source[key] = value
535 return source
537_ALLOWED_TOP_LEVEL_KEYS = frozenset({"model", "data", "train"})
540def _validate_top_level(config: Dict[str, Any]) -> None:
541 """Reject any top-level key other than ``model`` / ``data`` / ``train``.
543 Strict three-tier YAML shape. Any flat-style legacy key
544 (``parallel``, ``optim``, ``mixed_precision``, ``runtime``, ``debug`` ...)
545 must be moved under ``train.*`` — see schema.py for the canonical layout.
547 Raises:
548 ValueError: With migration hints when forbidden top-level keys are
549 present in the YAML.
550 """
551 forbidden = sorted(set(config) - _ALLOWED_TOP_LEVEL_KEYS)
552 if not forbidden:
553 return
555 legacy_to_train_path = {
556 "parallel": "train.accelerator",
557 "optim": "train.optimizer",
558 "mixed_precision": "train.mixed_precision",
559 "memory": "train.gradient_checkpointing",
560 "checkpoint": "train.checkpoint",
561 "logging": "train.logging",
562 "tensorboard": "train.tensorboard",
563 "wandb": "train.wandb",
564 "profiler": "train.profile",
565 "memory_monitor": "train.memory_monitor",
566 "moe_monitor": "train.moe_monitor",
567 "eval": "train.eval",
568 "runtime": "train (flatten init_device / backend / comm_backend)",
569 "debug": "train.debug",
570 "seed": "train.seed",
571 }
572 hints = []
573 for key in forbidden:
574 new_path = legacy_to_train_path.get(key)
575 if new_path:
576 hints.append(f" - top-level '{key}:' → move under {new_path}")
577 else:
578 hints.append(f" - top-level '{key}:' is not allowed")
579 raise ValueError(
580 "Forbidden top-level YAML keys: %s. The schema is strict three-tier "
581 "(model / data / train) — see config/schema.py. Migrate as follows:\n%s"
582 % (forbidden, "\n".join(hints))
583 )
586def _instantiate_recursive(cls: Type[T], config_dict: Dict[str, Any]) -> T:
587 """Recursively convert a dict into nested dataclass instances.
589 Unknown keys in ``config_dict`` that have no corresponding field on
590 ``cls`` emit a ``logger.warning`` with a closest-match suggestion from
591 ``difflib``, helping users catch typos in YAML configs.
592 """
593 if not is_dataclass(cls):
594 return config_dict
596 known = {f.name for f in fields(cls)}
597 unknown = set(config_dict) - known
598 for name in sorted(unknown):
599 matches = difflib.get_close_matches(name, known, n=1)
600 suggestion = f" Did you mean '{matches[0]}'?" if matches else ""
601 logger.warning(
602 "Unknown config key '%s' for %s ignored.%s",
603 name, cls.__name__, suggestion,
604 )
606 field_values = {}
607 for field_info in fields(cls):
608 if field_info.name not in config_dict:
609 continue
610 raw_value = config_dict[field_info.name]
611 field_type = field_info.type
613 # Unwrap Optional[X] → X
614 origin = get_origin(field_type)
615 if origin is Union:
616 unwrapped = [a for a in get_args(field_type) if a is not type(None)]
617 if len(unwrapped) == 1:
618 field_type = unwrapped[0]
620 if is_dataclass(field_type) and isinstance(raw_value, dict):
621 field_values[field_info.name] = _instantiate_recursive(field_type, raw_value)
622 elif field_type is bool and isinstance(raw_value, str):
623 field_values[field_info.name] = _string_to_bool(raw_value)
624 else:
625 field_values[field_info.name] = raw_value
627 return cls(**field_values)
630def parse_args(root_class: Type[T]) -> T:
631 """Parse training config from YAML file + CLI overrides.
633 Usage::
635 args = parse_args(HyperTrainerConfig)
637 The first positional argument is the YAML config file path.
638 CLI arguments use dot-path notation under the strict three-tier schema:
639 ``--train.accelerator.dp_shard=8 --train.optimizer.lr=3e-4``
641 Bool fields accept extended string aliases (``yes/no/on/off/y/n/t/f/1/0``).
642 Unknown YAML keys emit a warning with a closest-match suggestion.
644 Args:
645 root_class: The root config dataclass type.
647 Returns:
648 An instance of root_class populated from YAML + CLI.
649 """
650 parser = argparse.ArgumentParser(description="HyperParallel Trainer")
651 parser.add_argument("config_file", nargs="?", help="Path to YAML config file")
652 args, remaining = parser.parse_known_args()
654 # Load YAML
655 final_config: dict = {}
656 if args.config_file:
657 if not os.path.isfile(args.config_file):
658 logger.warning(
659 "Config file not found: %s (cwd=%s). Using all defaults.",
660 args.config_file, os.getcwd(),
661 )
662 else:
663 with open(args.config_file, encoding="utf-8") as f:
664 yaml_config = yaml.safe_load(f)
665 if yaml_config:
666 final_config = yaml_config
668 # Parse CLI dot-path overrides: --train.accelerator.dp=8 → nested dict
669 cli_config: dict = {}
670 for item in remaining:
671 if item.startswith("--") and "=" in item:
672 dot_key, raw_value = item[2:].split("=", 1)
673 coerced = _coerce_cli_value(raw_value, dot_key, root_class)
674 keys = dot_key.split(".")
675 current = cli_config
676 for k in keys[:-1]:
677 current = current.setdefault(k, {})
678 current[keys[-1]] = coerced
680 # CLI overrides YAML
681 final_config = _deep_update(final_config, cli_config)
683 # Strict three-tier validation — only model / data / train allowed.
684 _validate_top_level(final_config)
686 # local_rank from environment (torchrun sets it).
687 local_rank = int(os.environ.get("LOCAL_RANK", "0"))
688 final_config.setdefault("train", {})["local_rank"] = local_rank
690 return _instantiate_recursive(root_class, final_config)