Add files using upload-large-folder tool

.gitattributes

@@ -49,3 +49,4 @@ lr_sweep/hnet_xl_code_lr_2e-4/wandb/run-20260425_222011-khn25dwv/run-khn25dwv.wa
 pythia_1_4b_rerun/pythia_1_4b_lr_2e-5/wandb/run-20260425_215604-jhqe4qjw/run-jhqe4qjw.wandb filter=lfs diff=lfs merge=lfs -text
 generation_bench_results/runs.partial.jsonl filter=lfs diff=lfs merge=lfs -text
 generation_bench_results/results.json filter=lfs diff=lfs merge=lfs -text
+pythia_1b_rerun/pythia_1b_lr_2e-5/wandb/run-20260523_224742-l9w4bbvz/run-l9w4bbvz.wandb filter=lfs diff=lfs merge=lfs -text

pythia_1b_rerun/pythia_1b_lr_2e-5/.hydra/config.yaml

@@ -0,0 +1,50 @@
+model:
+  name: EleutherAI/pythia-1b
+  checkpoint_path: null
+  from_scratch: false
+training:
+  epochs: 3
+  batch_size: 1
+  eval_batch_size: 1
+  gradient_accumulation_steps: 8
+  lr: 2.0e-05
+  weight_decay: 0.1
+  betas:
+  - 0.9
+  - 0.95
+  eps: 1.0e-08
+  lr_scheduler: wsd
+  warmup_ratio: 0.1
+  decay_ratio: 0.2
+  warmup_steps: 100
+  min_lr_ratio: 0.1
+  max_grad_norm: 1.0
+  use_amp: true
+  gradient_checkpointing: true
+  resume: false
+  resume_checkpoint: null
+data:
+  path: ${oc.env:PROJECT_ROOT}/code_completion_exp/datasets/data_V4_full
+  max_context_len: 4096
+  max_target_len: 256
+  num_workers: 4
+  pin_memory: true
+  max_train_samples: null
+  max_val_samples: 2000
+logging:
+  log_interval: 10
+  save_interval: 12000
+  eval_interval: 12000
+  save_every_epoch: true
+tracking:
+  enabled: true
+  backend: wandb
+  project: code-completion_pythia-1b-rerun
+  run_name: pythia_1b_lr_2e-5
+  entity: null
+  base_url: https://wandb.platun0v.ru
+  local_dir: ${paths.output_dir}
+paths:
+  output_dir: outputs/pythia_1b_rerun/pythia_1b_lr_2e-5
+seed: 42
+device: cuda

pythia_1b_rerun/pythia_1b_lr_2e-5/.hydra/hydra.yaml

@@ -0,0 +1,160 @@
+hydra:
+  run:
+    dir: ${paths.output_dir}
+  sweep:
+    dir: outputs/multirun/${now:%Y-%m-%d}/${now:%H-%M-%S}
+    subdir: ${hydra.job.num}
+  launcher:
+    _target_: hydra._internal.core_plugins.basic_launcher.BasicLauncher
+  sweeper:
+    _target_: hydra._internal.core_plugins.basic_sweeper.BasicSweeper
+    max_batch_size: null
+    params: null
+  help:
+    app_name: ${hydra.job.name}
+    header: '${hydra.help.app_name} is powered by Hydra.
+
+      '
+    footer: 'Powered by Hydra (https://hydra.cc)
+
+      Use --hydra-help to view Hydra specific help
+
+      '
+    template: '${hydra.help.header}
+
+      == Configuration groups ==
+
+      Compose your configuration from those groups (group=option)
+
+
+      $APP_CONFIG_GROUPS
+
+
+      == Config ==
+
+      Override anything in the config (foo.bar=value)
+
+
+      $CONFIG
+
+
+      ${hydra.help.footer}
+
+      '
+  hydra_help:
+    template: 'Hydra (${hydra.runtime.version})
+
+      See https://hydra.cc for more info.
+
+
+      == Flags ==
+
+      $FLAGS_HELP
+
+
+      == Configuration groups ==
+
+      Compose your configuration from those groups (For example, append hydra/job_logging=disabled
+      to command line)
+
+
+      $HYDRA_CONFIG_GROUPS
+
+
+      Use ''--cfg hydra'' to Show the Hydra config.
+
+      '
+    hydra_help: ???
+  hydra_logging:
+    version: 1
+    formatters:
+      simple:
+        format: '[%(asctime)s][HYDRA] %(message)s'
+    handlers:
+      console:
+        class: logging.StreamHandler
+        formatter: simple
+        stream: ext://sys.stdout
+    root:
+      level: INFO
+      handlers:
+      - console
+    loggers:
+      logging_example:
+        level: DEBUG
+    disable_existing_loggers: false
+  job_logging:
+    version: 1
+    formatters:
+      simple:
+        format: '[%(asctime)s][%(name)s][%(levelname)s] - %(message)s'
+    handlers:
+      console:
+        class: logging.StreamHandler
+        formatter: simple
+        stream: ext://sys.stdout
+      file:
+        class: logging.FileHandler
+        formatter: simple
+        filename: ${hydra.runtime.output_dir}/${hydra.job.name}.log
+    root:
+      level: INFO
+      handlers:
+      - console
+      - file
+    disable_existing_loggers: false
+  env: {}
+  mode: RUN
+  searchpath: []
+  callbacks: {}
+  output_subdir: .hydra
+  overrides:
+    hydra:
+    - hydra.mode=RUN
+    task: []
+  job:
+    name: train
+    chdir: false
+    override_dirname: ''
+    id: ???
+    num: ???
+    config_name: config
+    env_set: {}
+    env_copy: []
+    config:
+      override_dirname:
+        kv_sep: '='
+        item_sep: ','
+        exclude_keys: []
+  runtime:
+    version: 1.3.2
+    version_base: '1.3'
+    cwd: /workspace/byte-llms-code/code_completion_exp/train_pythia
+    config_sources:
+    - path: hydra.conf
+      schema: pkg
+      provider: hydra
+    - path: /workspace/byte-llms-code/code_completion_exp/train_pythia/configs
+      schema: file
+      provider: main
+    - path: ''
+      schema: structured
+      provider: schema
+    output_dir: /workspace/byte-llms-code/code_completion_exp/train_pythia/outputs/pythia_1b_rerun/pythia_1b_lr_2e-5
+    choices:
+      paths: default
+      tracking: wandb
+      logging: default
+      data: default
+      training: default
+      model: pythia_1b
+      hydra/env: default
+      hydra/callbacks: null
+      hydra/job_logging: default
+      hydra/hydra_logging: default
+      hydra/hydra_help: default
+      hydra/help: default
+      hydra/sweeper: basic
+      hydra/launcher: basic
+      hydra/output: default
+  verbose: false

pythia_1b_rerun/pythia_1b_lr_2e-5/.hydra/overrides.yaml

@@ -0,0 +1 @@
+[]

pythia_1b_rerun/pythia_1b_lr_2e-5/eval_results/eval_config.yaml

@@ -0,0 +1,31 @@
+data:
+  path: /workspace/byte-llms-code/code_completion_exp/datasets/data_V4_full
+  max_context_len: 4096
+  max_target_len: 256
+  num_workers: 4
+  pin_memory: true
+  max_train_samples: null
+  max_val_samples: 2000
+model:
+  name: EleutherAI/pythia-1b
+  checkpoint_path: null
+  from_scratch: false
+paths:
+  checkpoints_dir: outputs/pythia_1b_rerun/pythia_1b_lr_2e-5
+  initial_checkpoint: auto
+  output_dir: outputs/pythia_1b_rerun/pythia_1b_lr_2e-5/eval_results
+evaluation:
+  batch_size: 2
+  max_samples: 5000
+  compute_perplexity: true
+  bleu_tokenize: none
+  save_predictions: true
+  use_amp: true
+generation:
+  max_new_tokens: 64
+  temperature: 0.1
+  top_k: 0
+  top_p: 1.0
+  do_sample: true
+seed: 42
+device: cuda

pythia_1b_rerun/pythia_1b_lr_2e-5/eval_results/metrics_checkpoint_latest.txt

@@ -0,0 +1,17 @@
+Checkpoint: checkpoint_latest.pt
+================================================================================
+
+exact_match: 0.3482
+token_accuracy: 0.33828347830186456
+bleu: 17.8853044827809
+perplexity: 2.905595600384796
+num_samples: 5000
+gen_wall_time_s: 159.0116421286948
+gen_samples_per_s: 31.444238503953628
+gen_time_per_sample_ms: 31.802328425738963
+gen_chars_per_s: 728.7831158061332
+gen_batch_mean_ms: 254.41862740591168
+gen_batch_p50_ms: 229.27742509637028
+gen_batch_p95_ms: 445.24853378534294
+gen_batch_max_ms: 1709.8218339961022
+gen_num_batches: 625

pythia_1b_rerun/pythia_1b_lr_2e-5/eval_results/metrics_checkpoint_step_108000.txt

@@ -0,0 +1,17 @@
+Checkpoint: checkpoint_step_108000.pt
+================================================================================
+
+exact_match: 0.3402
+token_accuracy: 0.33731702142705605
+bleu: 17.930937581273497
+perplexity: 2.9304769575371843
+num_samples: 5000
+gen_wall_time_s: 157.80592869932298
+gen_samples_per_s: 31.684487656523967
+gen_time_per_sample_ms: 31.561185739864598
+gen_chars_per_s: 734.1929479769733
+gen_batch_mean_ms: 252.48948591891676
+gen_batch_p50_ms: 224.54757196828723
+gen_batch_p95_ms: 453.54253498371685
+gen_batch_max_ms: 1574.9100289540365
+gen_num_batches: 625

pythia_1b_rerun/pythia_1b_lr_2e-5/eval_results/metrics_checkpoint_step_118647.txt

@@ -0,0 +1,17 @@
+Checkpoint: checkpoint_step_118647.pt
+================================================================================
+
+exact_match: 0.3464
+token_accuracy: 0.33543125191523465
+bleu: 18.058935206080903
+perplexity: 2.905685361730921
+num_samples: 5000
+gen_wall_time_s: 214.4719856326701
+gen_samples_per_s: 23.313068069242323
+gen_time_per_sample_ms: 42.89439712653402
+gen_chars_per_s: 537.5900244495003
+gen_batch_mean_ms: 171.5775885061361
+gen_batch_p50_ms: 150.72167199105024
+gen_batch_p95_ms: 329.12455603945995
+gen_batch_max_ms: 835.7596239075065
+gen_num_batches: 1250

pythia_1b_rerun/pythia_1b_lr_2e-5/eval_results/metrics_checkpoint_step_12000.txt

@@ -0,0 +1,17 @@
+Checkpoint: checkpoint_step_12000.pt
+================================================================================
+
+exact_match: 0.3108
+token_accuracy: 0.3025010018150531
+bleu: 16.718847641807816
+perplexity: 3.0961581194720997
+num_samples: 5000
+gen_wall_time_s: 202.44847195665352
+gen_samples_per_s: 24.697642573812836
+gen_time_per_sample_ms: 40.4896943913307
+gen_chars_per_s: 532.871396644071
+gen_batch_mean_ms: 161.9587775653228
+gen_batch_p50_ms: 137.1968500316143
+gen_batch_p95_ms: 327.32310149585817
+gen_batch_max_ms: 806.8903230596334
+gen_num_batches: 1250

pythia_1b_rerun/pythia_1b_lr_2e-5/eval_results/metrics_checkpoint_step_24000.txt

@@ -0,0 +1,17 @@
+Checkpoint: checkpoint_step_24000.pt
+================================================================================
+
+exact_match: 0.3166
+token_accuracy: 0.3113876906395116
+bleu: 16.67419761557939
+perplexity: 3.0413994225389254
+num_samples: 5000
+gen_wall_time_s: 208.3465970712714
+gen_samples_per_s: 23.998472114663794
+gen_time_per_sample_ms: 41.66931941425428
+gen_chars_per_s: 534.5371681875985
+gen_batch_mean_ms: 166.6772776570171
+gen_batch_p50_ms: 142.52355351345614
+gen_batch_p95_ms: 330.15160368522623
+gen_batch_max_ms: 825.7469990057871
+gen_num_batches: 1250

pythia_1b_rerun/pythia_1b_lr_2e-5/eval_results/metrics_checkpoint_step_36000.txt

@@ -0,0 +1,17 @@
+Checkpoint: checkpoint_step_36000.pt
+================================================================================
+
+exact_match: 0.3252
+token_accuracy: 0.321900855667916
+bleu: 18.101021072851456
+perplexity: 2.9839677869713896
+num_samples: 5000
+gen_wall_time_s: 222.01054317341186
+gen_samples_per_s: 22.521452938811617
+gen_time_per_sample_ms: 44.40210863468237
+gen_chars_per_s: 509.86767737234396
+gen_batch_mean_ms: 177.60843453872948
+gen_batch_p50_ms: 151.81013598339632
+gen_batch_p95_ms: 359.47523709037335
+gen_batch_max_ms: 1223.8604919984937
+gen_num_batches: 1250

pythia_1b_rerun/pythia_1b_lr_2e-5/eval_results/metrics_checkpoint_step_39549.txt

@@ -0,0 +1,17 @@
+Checkpoint: checkpoint_step_39549.pt
+================================================================================
+
+exact_match: 0.3276
+token_accuracy: 0.32032152370176553
+bleu: 17.54822744008585
+perplexity: 2.965576807077662
+num_samples: 5000
+gen_wall_time_s: 214.7674574032426
+gen_samples_per_s: 23.280994525218556
+gen_time_per_sample_ms: 42.953491480648516
+gen_chars_per_s: 541.0456565671742
+gen_batch_mean_ms: 171.81396592259406
+gen_batch_p50_ms: 147.50543044647202
+gen_batch_p95_ms: 333.53986238944344
+gen_batch_max_ms: 1203.0604279134423
+gen_num_batches: 1250

pythia_1b_rerun/pythia_1b_lr_2e-5/eval_results/metrics_checkpoint_step_48000.txt

@@ -0,0 +1,17 @@
+Checkpoint: checkpoint_step_48000.pt
+================================================================================
+
+exact_match: 0.3274
+token_accuracy: 0.32296160101831556
+bleu: 16.73652513019122
+perplexity: 2.9918445958286943
+num_samples: 5000
+gen_wall_time_s: 219.35748579574283
+gen_samples_per_s: 22.793842580124235
+gen_time_per_sample_ms: 43.87149715914857
+gen_chars_per_s: 532.0174033571317
+gen_batch_mean_ms: 175.48598863659427
+gen_batch_p50_ms: 153.4417055081576
+gen_batch_p95_ms: 343.1636362220159
+gen_batch_max_ms: 1041.408182005398
+gen_num_batches: 1250

pythia_1b_rerun/pythia_1b_lr_2e-5/eval_results/metrics_checkpoint_step_60000.txt

@@ -0,0 +1,17 @@
+Checkpoint: checkpoint_step_60000.pt
+================================================================================
+
+exact_match: 0.3282
+token_accuracy: 0.31928435047026377
+bleu: 18.72282952157622
+perplexity: 2.993910315820002
+num_samples: 5000
+gen_wall_time_s: 203.61096290219575
+gen_samples_per_s: 24.55663451875007
+gen_time_per_sample_ms: 40.72219258043915
+gen_chars_per_s: 549.4939879722629
+gen_batch_mean_ms: 162.8887703217566
+gen_batch_p50_ms: 143.41317699290812
+gen_batch_p95_ms: 310.04469925537705
+gen_batch_max_ms: 1205.2420629188418
+gen_num_batches: 1250

pythia_1b_rerun/pythia_1b_lr_2e-5/eval_results/metrics_checkpoint_step_72000.txt

@@ -0,0 +1,17 @@
+Checkpoint: checkpoint_step_72000.pt
+================================================================================
+
+exact_match: 0.322
+token_accuracy: 0.32331518280178206
+bleu: 16.71863598766819
+perplexity: 2.9653935325591534
+num_samples: 5000
+gen_wall_time_s: 309.09055828885175
+gen_samples_per_s: 16.176488947706364
+gen_time_per_sample_ms: 61.818111657770345
+gen_chars_per_s: 375.99983850490764
+gen_batch_mean_ms: 123.6362233155407
+gen_batch_p50_ms: 106.99726996244863
+gen_batch_p95_ms: 258.04170356132084
+gen_batch_max_ms: 788.3197760675102
+gen_num_batches: 2500

pythia_1b_rerun/pythia_1b_lr_2e-5/eval_results/metrics_checkpoint_step_79098.txt

@@ -0,0 +1,17 @@
+Checkpoint: checkpoint_step_79098.pt
+================================================================================
+
+exact_match: 0.3298
+token_accuracy: 0.320887254555312
+bleu: 17.71138353050089
+perplexity: 2.942476525271376
+num_samples: 5000
+gen_wall_time_s: 309.0561787857441
+gen_samples_per_s: 16.178288425245473
+gen_time_per_sample_ms: 61.81123575714882
+gen_chars_per_s: 372.56980414497804
+gen_batch_mean_ms: 123.62247151429764
+gen_batch_p50_ms: 107.23386402241886
+gen_batch_p95_ms: 263.92296123667614
+gen_batch_max_ms: 876.8716310150921
+gen_num_batches: 2500

pythia_1b_rerun/pythia_1b_lr_2e-5/eval_results/metrics_checkpoint_step_84000.txt

@@ -0,0 +1,17 @@
+Checkpoint: checkpoint_step_84000.pt
+================================================================================
+
+exact_match: 0.329
+token_accuracy: 0.3245880772222615
+bleu: 16.9617629224614
+perplexity: 3.0389716528850443
+num_samples: 5000
+gen_wall_time_s: 314.8586326778168
+gen_samples_per_s: 15.880142645211558
+gen_time_per_sample_ms: 62.97172653556336
+gen_chars_per_s: 370.6679375674701
+gen_batch_mean_ms: 125.94345307112671
+gen_batch_p50_ms: 107.21351049141958
+gen_batch_p95_ms: 267.4172037630342
+gen_batch_max_ms: 861.1630159430206
+gen_num_batches: 2500

pythia_1b_rerun/pythia_1b_lr_2e-5/eval_results/metrics_checkpoint_step_96000.txt

@@ -0,0 +1,17 @@
+Checkpoint: checkpoint_step_96000.pt
+================================================================================
+
+exact_match: 0.3256
+token_accuracy: 0.3261202649506164
+bleu: 17.503664376154067
+perplexity: 3.005788407729942
+num_samples: 5000
+gen_wall_time_s: 315.26597137225326
+gen_samples_per_s: 15.859624742361436
+gen_time_per_sample_ms: 63.05319427445065
+gen_chars_per_s: 372.6789779708544
+gen_batch_mean_ms: 126.10638854890131
+gen_batch_p50_ms: 107.08245052956045
+gen_batch_p95_ms: 268.73391066328605
+gen_batch_max_ms: 873.0702099855989
+gen_num_batches: 2500

pythia_1b_rerun/pythia_1b_lr_2e-5/eval_results/metrics_initial_checkpoint.txt

@@ -0,0 +1,17 @@
+Checkpoint: initial_checkpoint
+================================================================================
+
+exact_match: 0.0
+token_accuracy: 0.2584211394762275
+bleu: 0.9735244963491229
+perplexity: 12.08306519007595
+num_samples: 5000
+gen_wall_time_s: 505.01385805231985
+gen_samples_per_s: 9.900718406586767
+gen_time_per_sample_ms: 101.00277161046398
+gen_chars_per_s: 2175.0591245878272
+gen_batch_mean_ms: 808.0221728837117
+gen_batch_p50_ms: 783.4902260219678
+gen_batch_p95_ms: 907.4296850711106
+gen_batch_max_ms: 1885.9824809478596
+gen_num_batches: 625

pythia_1b_rerun/pythia_1b_lr_2e-5/eval_results/metrics_model_best.txt

@@ -0,0 +1,17 @@
+Checkpoint: model_best.pt
+================================================================================
+
+exact_match: 0.342
+token_accuracy: 0.33460622775381277
+bleu: 18.229909693407816
+perplexity: 2.930694069292247
+num_samples: 5000
+gen_wall_time_s: 315.6722708721645
+gen_samples_per_s: 15.839211933900945
+gen_time_per_sample_ms: 63.1344541744329
+gen_chars_per_s: 366.76645585663675
+gen_batch_mean_ms: 126.2689083488658
+gen_batch_p50_ms: 108.94419148098677
+gen_batch_p95_ms: 268.3883103250988
+gen_batch_max_ms: 867.4433380365372
+gen_num_batches: 2500

pythia_1b_rerun/pythia_1b_lr_2e-5/eval_results/metrics_model_final.txt

@@ -0,0 +1,17 @@
+Checkpoint: model_final.pt
+================================================================================
+
+exact_match: 0.3458
+token_accuracy: 0.3345119392782217
+bleu: 17.63206076668849
+perplexity: 2.9054952795351836
+num_samples: 5000
+gen_wall_time_s: 315.89769689075183
+gen_samples_per_s: 15.82790900096106
+gen_time_per_sample_ms: 63.179539378150366
+gen_chars_per_s: 365.5075713955934
+gen_batch_mean_ms: 126.35907875630073
+gen_batch_p50_ms: 108.14683401258662
+gen_batch_p95_ms: 269.0969737188425
+gen_batch_max_ms: 891.0078270127997
+gen_num_batches: 2500

pythia_1b_rerun/pythia_1b_lr_2e-5/eval_results/summary.txt

@@ -0,0 +1,11 @@
+EVALUATION SUMMARY
+==================================================================================================
+
+Checkpoint                      Exact Match    Token Acc     BLEU   PERPLEXITY  ms/sample   samp/s
+--------------------------------------------------------------------------------------------------
+checkpoint_step_72000                32.20%       32.33%    16.72         2.97       61.8    16.18
+checkpoint_step_79098                32.98%       32.09%    17.71         2.94       61.8    16.18
+checkpoint_step_84000                32.90%       32.46%    16.96         3.04       63.0    15.88
+checkpoint_step_96000                32.56%       32.61%    17.50         3.01       63.1    15.86
+model_best                           34.20%       33.46%    18.23         2.93       63.1    15.84
+model_final                          34.58%       33.45%    17.63         2.91       63.2    15.83

pythia_1b_rerun/pythia_1b_lr_2e-5/wandb/debug-internal.log

@@ -0,0 +1,13 @@
+{"time":"2026-05-23T22:47:42.958656706Z","level":"INFO","msg":"stream: starting","core version":"0.24.0"}
+{"time":"2026-05-23T22:47:43.340677291Z","level":"INFO","msg":"stream: created new stream","id":"l9w4bbvz"}
+{"time":"2026-05-23T22:47:43.340753386Z","level":"INFO","msg":"handler: started","stream_id":"l9w4bbvz"}
+{"time":"2026-05-23T22:47:43.340906156Z","level":"INFO","msg":"stream: started","id":"l9w4bbvz"}
+{"time":"2026-05-23T22:47:43.340930302Z","level":"INFO","msg":"writer: started","stream_id":"l9w4bbvz"}
+{"time":"2026-05-23T22:47:43.340957002Z","level":"INFO","msg":"sender: started","stream_id":"l9w4bbvz"}
+{"time":"2026-05-23T22:47:43.525566664Z","level":"ERROR","msg":"git repo not found","error":"repository does not exist"}
+{"time":"2026-05-24T15:56:08.713224566Z","level":"INFO","msg":"fileTransfer: Close: file transfer manager closed"}
+{"time":"2026-05-24T15:56:08.816433414Z","level":"INFO","msg":"handler: operation stats","stats":{}}
+{"time":"2026-05-24T15:56:08.820223515Z","level":"INFO","msg":"stream: closing","id":"l9w4bbvz"}
+{"time":"2026-05-24T15:56:08.820246589Z","level":"INFO","msg":"handler: closed","stream_id":"l9w4bbvz"}
+{"time":"2026-05-24T15:56:08.820433113Z","level":"INFO","msg":"sender: closed","stream_id":"l9w4bbvz"}
+{"time":"2026-05-24T15:56:08.820445777Z","level":"INFO","msg":"stream: closed","id":"l9w4bbvz"}

pythia_1b_rerun/pythia_1b_lr_2e-5/wandb/debug.log

@@ -0,0 +1,24 @@
+2026-05-23 22:47:42,622 INFO    MainThread:31611 [wandb_setup.py:_flush():81] Current SDK version is 0.24.0
+2026-05-23 22:47:42,623 INFO    MainThread:31611 [wandb_setup.py:_flush():81] Configure stats pid to 31611
+2026-05-23 22:47:42,623 INFO    MainThread:31611 [wandb_setup.py:_flush():81] Loading settings from environment variables
+2026-05-23 22:47:42,623 INFO    MainThread:31611 [wandb_init.py:setup_run_log_directory():717] Logging user logs to outputs/pythia_1b_rerun/pythia_1b_lr_2e-5/wandb/run-20260523_224742-l9w4bbvz/logs/debug.log
+2026-05-23 22:47:42,623 INFO    MainThread:31611 [wandb_init.py:setup_run_log_directory():718] Logging internal logs to outputs/pythia_1b_rerun/pythia_1b_lr_2e-5/wandb/run-20260523_224742-l9w4bbvz/logs/debug-internal.log
+2026-05-23 22:47:42,623 INFO    MainThread:31611 [wandb_init.py:init():844] calling init triggers
+2026-05-23 22:47:42,623 INFO    MainThread:31611 [wandb_init.py:init():849] wandb.init called with sweep_config: {}
+config: {'model': {'name': 'EleutherAI/pythia-1b', 'checkpoint_path': None, 'from_scratch': False}, 'training': {'epochs': 3, 'batch_size': 1, 'eval_batch_size': 1, 'gradient_accumulation_steps': 8, 'lr': 2e-05, 'weight_decay': 0.1, 'betas': [0.9, 0.95], 'eps': 1e-08, 'lr_scheduler': 'wsd', 'warmup_ratio': 0.1, 'decay_ratio': 0.2, 'warmup_steps': 100, 'min_lr_ratio': 0.1, 'max_grad_norm': 1.0, 'use_amp': True, 'gradient_checkpointing': True, 'resume': False, 'resume_checkpoint': None}, 'data': {'path': '/workspace/byte-llms-code/code_completion_exp/datasets/data_V4_full', 'max_context_len': 4096, 'max_target_len': 256, 'num_workers': 4, 'pin_memory': True, 'max_train_samples': None, 'max_val_samples': 2000}, 'logging': {'log_interval': 10, 'save_interval': 12000, 'eval_interval': 12000, 'save_every_epoch': True}, 'tracking': {'enabled': True, 'backend': 'wandb', 'project': 'code-completion_pythia-1b-rerun', 'run_name': 'pythia_1b_lr_2e-5', 'entity': None, 'base_url': 'https://wandb.platun0v.ru', 'local_dir': 'outputs/pythia_1b_rerun/pythia_1b_lr_2e-5'}, 'paths': {'output_dir': 'outputs/pythia_1b_rerun/pythia_1b_lr_2e-5'}, 'seed': 42, 'device': 'cuda', '_wandb': {'code_path': 'code/code_completion_exp/train_pythia/train.py'}}
+2026-05-23 22:47:42,623 INFO    MainThread:31611 [wandb_init.py:init():892] starting backend
+2026-05-23 22:47:42,921 INFO    MainThread:31611 [wandb_init.py:init():895] sending inform_init request
+2026-05-23 22:47:42,957 INFO    MainThread:31611 [wandb_init.py:init():903] backend started and connected
+2026-05-23 22:47:42,960 INFO    MainThread:31611 [wandb_init.py:init():973] updated telemetry
+2026-05-23 22:47:42,996 INFO    MainThread:31611 [wandb_init.py:init():997] communicating run to backend with 90.0 second timeout
+2026-05-23 22:47:43,524 INFO    MainThread:31611 [wandb_init.py:init():1044] starting run threads in backend
+2026-05-23 22:47:43,678 INFO    MainThread:31611 [wandb_run.py:_console_start():2529] atexit reg
+2026-05-23 22:47:43,678 INFO    MainThread:31611 [wandb_run.py:_redirect():2377] redirect: wrap_raw
+2026-05-23 22:47:43,678 INFO    MainThread:31611 [wandb_run.py:_redirect():2446] Wrapping output streams.
+2026-05-23 22:47:43,678 INFO    MainThread:31611 [wandb_run.py:_redirect():2469] Redirects installed.
+2026-05-23 22:47:43,681 INFO    MainThread:31611 [wandb_init.py:init():1084] run started, returning control to user process
+2026-05-24 15:56:07,662 INFO    MainThread:31611 [wandb_run.py:_finish():2295] finishing run nikita/code-completion_pythia-1b-rerun/l9w4bbvz
+2026-05-24 15:56:07,662 INFO    MainThread:31611 [wandb_run.py:_atexit_cleanup():2494] got exitcode: 0
+2026-05-24 15:56:07,662 INFO    MainThread:31611 [wandb_run.py:_restore():2476] restore
+2026-05-24 15:56:07,662 INFO    MainThread:31611 [wandb_run.py:_restore():2482] restore done
+2026-05-24 15:56:08,819 INFO    MainThread:31611 [wandb_run.py:_footer_sync_info():3870] logging synced files

pythia_1b_rerun/pythia_1b_lr_2e-5/wandb/run-20260523_224742-l9w4bbvz/files/code/code_completion_exp/train_pythia/train.py

@@ -0,0 +1,615 @@
+"""
+Training Pipeline для Pythia (decoder-only transformer) на задаче Code Completion.
+
+Конфигурация через Hydra + OmegaConf, логирование в Trackio.
+Поддержка DDP через Accelerate для multi-GPU тренировки.
+
+Использование:
+    # Базовый запуск (single GPU)
+    python train.py
+
+    # Multi-GPU с Accelerate
+    accelerate launch train.py
+
+    # Multi-GPU с указанием количества GPU
+    accelerate launch --num_processes=4 train.py
+
+    # Переопределение параметров через CLI
+    python train.py training.lr=1e-4 training.epochs=5
+
+    # Выбор другого конфига модели
+    python train.py model=pythia_160m
+
+    # Multirun (sweep)
+    python train.py --multirun training.lr=1e-4,3e-4,1e-3
+
+    # Без логирования
+    python train.py tracking.enabled=false
+"""
+
+import os
+import math
+import time
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from datasets import load_from_disk
+
+import hydra
+from hydra.core.hydra_config import HydraConfig
+from omegaconf import DictConfig, OmegaConf
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    AutoConfig,
+    PreTrainedTokenizerBase,
+)
+from accelerate import Accelerator
+from accelerate.utils import set_seed as accelerate_set_seed
+
+# Ensure repo root is on sys.path (needed when running from subdirectory)
+import sys
+sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
+
+# Shared training library
+from training_lib.utils import AverageMeter, log_message
+from training_lib.checkpointing import save_checkpoint, load_checkpoint
+from training_lib.schedulers import get_lr_scheduler
+from training_lib.tracking import init_tracking, log_metrics, finish_tracking
+from training_lib.validation import run_validation
+
+
+# ============================================================================
+# ДАННЫЕ
+# ============================================================================
+
+
+class CodeCompletionCollator:
+    """Collate function для батчирования примеров code completion."""
+
+    def __init__(
+        self,
+        tokenizer: PreTrainedTokenizerBase,
+        max_context_len: int = 1024,
+        max_target_len: int = 256,
+    ):
+        self.tokenizer = tokenizer
+        self.max_context_len = max_context_len
+        self.max_target_len = max_target_len
+        self.pad_token_id = tokenizer.pad_token_id
+
+    def __call__(self, batch: list[dict]) -> dict:
+        contexts = [item["context"] for item in batch]
+        targets = [item["target"] for item in batch]
+
+        encoded_contexts = self.tokenizer(
+            contexts,
+            add_special_tokens=True,
+            truncation=True,
+            max_length=self.max_context_len,
+            return_tensors=None,
+        )
+        encoded_targets = self.tokenizer(
+            targets,
+            add_special_tokens=False,
+            truncation=True,
+            max_length=self.max_target_len,
+            return_tensors=None,
+        )
+
+        input_ids_list = []
+        context_lengths = []
+
+        for ctx_ids, tgt_ids in zip(
+            encoded_contexts["input_ids"], encoded_targets["input_ids"]
+        ):
+            tgt_ids = tgt_ids + [self.tokenizer.eos_token_id]
+            context_lengths.append(len(ctx_ids))
+            input_ids_list.append(ctx_ids + tgt_ids)
+
+        max_len = max(len(ids) for ids in input_ids_list)
+
+        padded_input_ids = []
+        attention_mask = []
+
+        for ids in input_ids_list:
+            padding_len = max_len - len(ids)
+            padded_input_ids.append(ids + [self.pad_token_id] * padding_len)
+            attention_mask.append([1] * len(ids) + [0] * padding_len)
+
+        return {
+            "input_ids": torch.tensor(padded_input_ids, dtype=torch.long),
+            "attention_mask": torch.tensor(attention_mask, dtype=torch.long),
+            "context_lengths": torch.tensor(context_lengths, dtype=torch.long),
+        }
+
+
+def create_dataloaders(
+    cfg: DictConfig, tokenizer: PreTrainedTokenizerBase
+) -> dict[str, DataLoader]:
+    """Создание DataLoader'ов для train и validation."""
+    dataset_dict = load_from_disk(cfg.data.path)
+
+    collator = CodeCompletionCollator(
+        tokenizer=tokenizer,
+        max_context_len=cfg.data.max_context_len,
+        max_target_len=cfg.data.max_target_len,
+    )
+
+    dataloaders = {}
+
+    if "train" in dataset_dict:
+        train_dataset = dataset_dict["train"]
+        max_train = cfg.data.get("max_train_samples", None)
+        if max_train is not None:
+            train_dataset = train_dataset.select(range(min(max_train, len(train_dataset))))
+        dataloaders["train"] = DataLoader(
+            train_dataset,
+            batch_size=cfg.training.batch_size,
+            shuffle=True,
+            collate_fn=collator,
+            num_workers=cfg.data.num_workers,
+            pin_memory=cfg.data.pin_memory,
+        )
+
+    if "validation" in dataset_dict:
+        val_dataset = dataset_dict["validation"]
+        max_val = cfg.data.get("max_val_samples", None)
+        if max_val is not None:
+            val_dataset = val_dataset.select(range(min(max_val, len(val_dataset))))
+        eval_batch_size = cfg.training.get("eval_batch_size", cfg.training.batch_size)
+        dataloaders["validation"] = DataLoader(
+            val_dataset,
+            batch_size=eval_batch_size,
+            shuffle=False,
+            collate_fn=collator,
+            num_workers=cfg.data.num_workers,
+            pin_memory=cfg.data.pin_memory,
+        )
+
+    return dataloaders
+
+
+
+
+# ============================================================================
+# LOSS ФУНКЦИИ
+# ============================================================================
+
+
+def compute_loss(
+    logits: torch.Tensor,
+    input_ids: torch.Tensor,
+    context_lengths: torch.Tensor,
+    attention_mask: torch.Tensor,
+) -> dict:
+    """Вычисление loss для авторегрессионной модели."""
+    batch_size, seq_len, vocab_size = logits.shape
+
+    shift_logits = logits[:, :-1, :].contiguous()
+    shift_labels = input_ids[:, 1:].contiguous()
+    shift_mask = attention_mask[:, 1:].contiguous()
+
+    target_mask = torch.zeros_like(shift_labels, dtype=torch.bool)
+    for i in range(batch_size):
+        ctx_len = context_lengths[i].item()
+        target_mask[i, ctx_len - 1 :] = True
+
+    final_mask = target_mask & shift_mask.bool()
+
+    if final_mask.sum() > 0:
+        loss = F.cross_entropy(
+            shift_logits[final_mask], shift_labels[final_mask], reduction="mean"
+        )
+    else:
+        loss = torch.tensor(0.0, device=logits.device)
+
+    return {"loss": loss}
+
+
+def _pythia_forward_loss(
+    model: nn.Module,
+    batch: dict,
+    cfg: DictConfig,
+    accelerator: Accelerator,
+) -> dict:
+    """Forward + loss for a plain HF causal LM (attention_mask= kwarg, .logits)."""
+    input_ids = batch["input_ids"]
+    attention_mask = batch["attention_mask"]
+    context_lengths = batch["context_lengths"]
+    output = model(input_ids, attention_mask=attention_mask)
+    return compute_loss(output.logits, input_ids, context_lengths, attention_mask)
+
+
+# ============================================================================
+# PARAMETER GROUPING
+# ============================================================================
+
+
+def group_params(model: nn.Module, weight_decay: float) -> list[dict]:
+    """Группировка параметров для optimizer."""
+    decay_params = []
+    no_decay_params = []
+    
+    for name, param in model.named_parameters():
+        if not param.requires_grad:
+            continue
+            
+        if "bias" in name or "LayerNorm" in name or "layernorm" in name:
+            no_decay_params.append(param)
+        else:
+            decay_params.append(param)
+    
+    return [
+        {"params": decay_params, "weight_decay": weight_decay},
+        {"params": no_decay_params, "weight_decay": 0.0},
+    ]
+
+
+
+
+# ============================================================================
+# TRAINING LOOP
+# ============================================================================
+
+
+def train_epoch(
+    model: nn.Module,
+    dataloader: DataLoader,
+    optimizer: torch.optim.Optimizer,
+    scheduler,
+    cfg: DictConfig,
+    epoch: int,
+    global_step: int,
+    accelerator: Accelerator,
+    val_dataloader: DataLoader | None = None,
+    best_val_loss: float = float("inf"),
+) -> tuple[int, float]:
+    """Один epoch тренировки. Возвращает (global_step, best_val_loss)."""
+    model.train()
+
+    loss_meter = AverageMeter()
+
+    optimizer.zero_grad()
+    accumulated_loss = 0.0
+    accumulated_steps = 0
+
+    epoch_start_time = time.time()
+    step_start_time = time.time()
+
+    for batch_idx, batch in enumerate(dataloader):
+        input_ids = batch["input_ids"]
+        attention_mask = batch["attention_mask"]
+        context_lengths = batch["context_lengths"]
+
+        with accelerator.autocast():
+            output = model(input_ids, attention_mask=attention_mask)
+            logits = output.logits
+            loss_dict = compute_loss(
+                logits, input_ids, context_lengths, attention_mask
+            )
+
+        loss = loss_dict["loss"] / cfg.training.gradient_accumulation_steps
+        accelerator.backward(loss)
+
+        accumulated_loss += loss_dict["loss"].item()
+        accumulated_steps += 1
+
+        if accumulated_steps == cfg.training.gradient_accumulation_steps:
+            if cfg.training.max_grad_norm > 0:
+                accelerator.clip_grad_norm_(
+                    model.parameters(), cfg.training.max_grad_norm
+                )
+
+            optimizer.step()
+            scheduler.step()
+            optimizer.zero_grad()
+
+            avg_loss = accumulated_loss / cfg.training.gradient_accumulation_steps
+            loss_meter.update(avg_loss)
+
+            global_step += 1
+
+            if global_step % cfg.logging.log_interval == 0:
+                step_time = time.time() - step_start_time
+                current_lr = scheduler.get_last_lr()[0]
+
+                metrics = {
+                    "train/loss": loss_meter.val,
+                    "train/loss_avg": loss_meter.avg,
+                    "train/lr": current_lr,
+                    "train/epoch": epoch,
+                    "train/step_time": step_time / cfg.logging.log_interval,
+                }
+
+                log_metrics(metrics, step=global_step)
+
+                log_message(
+                    f"Epoch {epoch} | Step {global_step} | "
+                    f"Loss: {loss_meter.avg:.4f} | "
+                    f"LR: {current_lr:.2e}",
+                    cfg,
+                    accelerator,
+                )
+
+                step_start_time = time.time()
+
+            if (
+                cfg.logging.save_interval > 0
+                and global_step % cfg.logging.save_interval == 0
+            ):
+                save_checkpoint(
+                    model, optimizer, scheduler, global_step, epoch, cfg, accelerator
+                )
+
+            eval_interval = cfg.logging.get("eval_interval", 0)
+            if (
+                eval_interval > 0
+                and val_dataloader is not None
+                and global_step % eval_interval == 0
+            ):
+                val_metrics = run_validation(
+                    model=model,
+                    dataloader=val_dataloader,
+                    cfg=cfg,
+                    global_step=global_step,
+                    accelerator=accelerator,
+                    forward_loss_fn=_pythia_forward_loss,
+                )
+
+                if val_metrics["val/loss"] < best_val_loss:
+                    best_val_loss = val_metrics["val/loss"]
+                    if accelerator.is_main_process:
+                        best_model_path = Path(cfg.paths.output_dir) / "model_best.pt"
+                        unwrapped_model = accelerator.unwrap_model(model)
+                        torch.save(unwrapped_model.state_dict(), best_model_path)
+                        log_message(
+                            f"New best model saved! Val loss: {best_val_loss:.4f}", 
+                            cfg, 
+                            accelerator
+                        )
+
+                        log_metrics(
+                            {
+                                "best/val_loss": best_val_loss,
+                                "best/val_perplexity": val_metrics["val/perplexity"],
+                                "best/step": global_step,
+                            },
+                            step=global_step,
+                        )
+
+                model.train()
+
+            accumulated_loss = 0.0
+            accumulated_steps = 0
+
+    epoch_time = time.time() - epoch_start_time
+    
+    log_message(
+        f"Epoch {epoch} completed in {epoch_time:.2f}s | "
+        f"Loss: {loss_meter.avg:.4f}",
+        cfg,
+        accelerator,
+    )
+
+    log_metrics({
+        "epoch/loss": loss_meter.avg,
+        "epoch/time": epoch_time,
+    })
+
+    return global_step, best_val_loss
+
+
+# ============================================================================
+# MAIN
+# ============================================================================
+
+
+@hydra.main(version_base=None, config_path="configs", config_name="config")
+def main(cfg: DictConfig):
+    """Главная функция тренировки с поддержкой DDP через Accelerate."""
+
+    # === Performance: Enable TF32 for faster matmuls on Ampere+ GPUs ===
+    torch.set_float32_matmul_precision('high')
+
+    # === Accelerator Setup ===
+    mixed_precision = "bf16" if cfg.training.use_amp else "no"
+    
+    accelerator = Accelerator(
+        mixed_precision=mixed_precision,
+        gradient_accumulation_steps=cfg.training.gradient_accumulation_steps,
+    )
+
+    # === Setup ===
+    accelerate_set_seed(cfg.seed)
+    
+    if cfg.paths.output_dir is None:
+        cfg.paths.output_dir = HydraConfig.get().runtime.output_dir
+
+    OmegaConf.resolve(cfg)
+
+    log_message(f"CUDA_VISIBLE_DEVICES: {os.environ.get('CUDA_VISIBLE_DEVICES', 'not set')}", cfg, accelerator)
+    log_message(f"Number of processes: {accelerator.num_processes}", cfg, accelerator)
+    log_message(f"Process index: {accelerator.process_index}", cfg, accelerator)
+    log_message(f"Mixed precision: {mixed_precision}", cfg, accelerator)
+    
+    log_message("=" * 60, cfg, accelerator)
+    log_message("Pythia Training Pipeline (Hydra + Trackio + Accelerate)", cfg, accelerator)
+    log_message("=" * 60, cfg, accelerator)
+    log_message(f"Config:\n{OmegaConf.to_yaml(cfg)}", cfg, accelerator)
+
+    # === Trackio Init ===
+    init_tracking(cfg, accelerator)
+
+    # === Tokenizer ===
+    log_message("Initializing tokenizer...", cfg, accelerator)
+    tokenizer = AutoTokenizer.from_pretrained(cfg.model.name)
+    
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+        tokenizer.pad_token_id = tokenizer.eos_token_id
+
+    # === Model ===
+    log_message("Loading model...", cfg, accelerator)
+    
+    # Flash Attention 2
+    torch_dtype = torch.bfloat16 if cfg.training.use_amp else torch.float32
+    
+    if cfg.model.checkpoint_path:
+        model = AutoModelForCausalLM.from_pretrained(
+            cfg.model.name,
+            attn_implementation="flash_attention_2",
+            torch_dtype=torch_dtype,
+        )
+        checkpoint = torch.load(cfg.model.checkpoint_path, map_location="cpu")
+        model.load_state_dict(checkpoint["model_state_dict"] if "model_state_dict" in checkpoint else checkpoint)
+        log_message(f"Loaded checkpoint: {cfg.model.checkpoint_path}", cfg, accelerator)
+    elif cfg.model.from_scratch:
+        config = AutoConfig.from_pretrained(cfg.model.name)
+        config._attn_implementation = "flash_attention_2"
+        model = AutoModelForCausalLM.from_config(config, torch_dtype=torch_dtype)
+        log_message(f"Initialized from scratch: {cfg.model.name}", cfg, accelerator)
+    else:
+        model = AutoModelForCausalLM.from_pretrained(
+            cfg.model.name,
+            attn_implementation="flash_attention_2",
+            torch_dtype=torch_dtype,
+        )
+        log_message(f"Loaded pretrained: {cfg.model.name}", cfg, accelerator)
+
+    if cfg.training.get("gradient_checkpointing", False):
+        if hasattr(model, "gradient_checkpointing_enable"):
+            model.gradient_checkpointing_enable()
+            if hasattr(model.config, "use_cache"):
+                model.config.use_cache = False
+            log_message("Gradient checkpointing enabled", cfg, accelerator)
+        else:
+            log_message("Gradient checkpointing requested but not supported by this model", cfg, accelerator)
+
+    model.train()
+
+    # Log model info
+    total_params = sum(p.numel() for p in model.parameters())
+    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    log_message(f"Total params: {total_params:,}", cfg, accelerator)
+    log_message(f"Trainable params: {trainable_params:,}", cfg, accelerator)
+
+    # === Data ===
+    log_message("Creating dataloaders...", cfg, accelerator)
+    dataloaders = create_dataloaders(cfg, tokenizer)
+
+    train_dataloader = dataloaders["train"]
+    val_dataloader = dataloaders.get("validation", None)
+
+    log_message(f"Train dataset size: {len(train_dataloader.dataset)}", cfg, accelerator)
+    log_message(f"Train batches per epoch (before DDP split): {len(train_dataloader)}", cfg, accelerator)
+
+    if val_dataloader:
+        log_message(f"Validation dataset size: {len(val_dataloader.dataset)}", cfg, accelerator)
+        log_message(f"Validation batches: {len(val_dataloader)}", cfg, accelerator)
+    else:
+        log_message("No validation dataset found", cfg, accelerator)
+
+    # === Optimizer ===
+    log_message("Creating optimizer...", cfg, accelerator)
+    param_groups = group_params(model, cfg.training.weight_decay)
+
+    optimizer = torch.optim.AdamW(
+        param_groups,
+        lr=cfg.training.lr,
+        betas=tuple(cfg.training.betas),
+        eps=cfg.training.eps,
+    )
+
+    # === Scheduler ===
+    steps_per_epoch = math.ceil(
+        len(train_dataloader) / accelerator.num_processes
+    )
+    total_steps = (
+        cfg.training.epochs
+        * steps_per_epoch
+        // cfg.training.gradient_accumulation_steps
+    )
+    scheduler = get_lr_scheduler(optimizer, cfg, total_steps)
+
+    log_message(
+        f"Total steps: {total_steps}, Steps per epoch: {steps_per_epoch}", 
+        cfg, 
+        accelerator
+    )
+
+    # === Accelerate Prepare ===
+    log_message("Preparing model, optimizer, and dataloaders with Accelerate...", cfg, accelerator)
+    
+    if val_dataloader is not None:
+        model, optimizer, train_dataloader, val_dataloader, scheduler = accelerator.prepare(
+            model, optimizer, train_dataloader, val_dataloader, scheduler
+        )
+    else:
+        model, optimizer, train_dataloader, scheduler = accelerator.prepare(
+            model, optimizer, train_dataloader, scheduler
+        )
+
+    log_message(f"Train batches per epoch (after DDP split): {len(train_dataloader)}", cfg, accelerator)
+
+    # === Resume ===
+    global_step = 0
+    start_epoch = 1
+
+    if cfg.training.resume and cfg.training.resume_checkpoint:
+        global_step, start_epoch = load_checkpoint(
+            model, optimizer, scheduler, cfg.training.resume_checkpoint, cfg, accelerator
+        )
+        start_epoch += 1
+
+    # === Training Loop ===
+    log_message("Starting training...", cfg, accelerator)
+
+    best_val_loss = float("inf")
+
+    try:
+        for epoch in range(start_epoch, cfg.training.epochs + 1):
+            log_message(f"\n{'=' * 60}", cfg, accelerator)
+            log_message(f"EPOCH {epoch}/{cfg.training.epochs}", cfg, accelerator)
+            log_message(f"{'=' * 60}", cfg, accelerator)
+
+            global_step, best_val_loss = train_epoch(
+                model=model,
+                dataloader=train_dataloader,
+                optimizer=optimizer,
+                scheduler=scheduler,
+                cfg=cfg,
+                epoch=epoch,
+                global_step=global_step,
+                accelerator=accelerator,
+                val_dataloader=val_dataloader,
+                best_val_loss=best_val_loss,
+            )
+
+            if cfg.logging.save_every_epoch:
+                save_checkpoint(
+                    model, optimizer, scheduler, global_step, epoch, cfg, accelerator
+                )
+
+    except KeyboardInterrupt:
+        log_message("Training interrupted by user", cfg, accelerator)
+        save_checkpoint(model, optimizer, scheduler, global_step, epoch, cfg, accelerator)
+
+    # === Final Save ===
+    log_message("\nTraining completed!", cfg, accelerator)
+
+    if accelerator.is_main_process:
+        final_model_path = Path(cfg.paths.output_dir) / "model_final.pt"
+        unwrapped_model = accelerator.unwrap_model(model)
+        torch.save(unwrapped_model.state_dict(), final_model_path)
+        log_message(f"Final model: {final_model_path}", cfg, accelerator)
+
+    accelerator.wait_for_everyone()
+    finish_tracking()
+
+
+if __name__ == "__main__":
+    main()

pythia_1b_rerun/pythia_1b_lr_2e-5/wandb/run-20260523_224742-l9w4bbvz/files/config.yaml

@@ -0,0 +1,124 @@
+_wandb:
+    value:
+        cli_version: 0.24.0
+        code_path: code/code_completion_exp/train_pythia/train.py
+        e:
+            j28hszir4t4f3z6gdnpmywfj4lyk0mju:
+                codePath: code_completion_exp/train_pythia/train.py
+                codePathLocal: train.py
+                cpu_count: 32
+                cpu_count_logical: 64
+                cudaVersion: "12.7"
+                disk:
+                    /:
+                        total: "207232172032"
+                        used: "44856893440"
+                email: nikita@local.ru
+                executable: /venv/bytellm/bin/python
+                git:
+                    commit: 71d0b08221a685157b26484b9c9ef2e9ced3e783
+                    remote: https://github.com/naryst/byte-llms-code.git
+                gpu: NVIDIA GeForce RTX 4090
+                gpu_count: 1
+                gpu_nvidia:
+                    - architecture: Ada
+                      cudaCores: 16384
+                      memoryTotal: "25757220864"
+                      name: NVIDIA GeForce RTX 4090
+                      uuid: GPU-f251e903-dbf7-df97-a438-a37f8506d6f9
+                host: c46590d05c20
+                memory:
+                    total: "270057459712"
+                os: Linux-5.15.0-177-generic-x86_64-with-glibc2.39
+                program: /workspace/byte-llms-code/code_completion_exp/train_pythia/train.py
+                python: CPython 3.12.0
+                root: outputs/pythia_1b_rerun/pythia_1b_lr_2e-5
+                startedAt: "2026-05-23T22:47:42.621546Z"
+                writerId: j28hszir4t4f3z6gdnpmywfj4lyk0mju
+        m: []
+        python_version: 3.12.0
+        t:
+            "1":
+                - 1
+                - 11
+                - 49
+                - 50
+                - 51
+                - 71
+                - 105
+            "2":
+                - 1
+                - 11
+                - 49
+                - 50
+                - 51
+                - 71
+                - 105
+            "3":
+                - 2
+                - 13
+                - 16
+                - 61
+            "4": 3.12.0
+            "5": 0.24.0
+            "6": 4.57.6
+            "12": 0.24.0
+            "13": linux-x86_64
+data:
+    value:
+        max_context_len: 4096
+        max_target_len: 256
+        max_train_samples: null
+        max_val_samples: 2000
+        num_workers: 4
+        path: /workspace/byte-llms-code/code_completion_exp/datasets/data_V4_full
+        pin_memory: true
+device:
+    value: cuda
+logging:
+    value:
+        eval_interval: 12000
+        log_interval: 10
+        save_every_epoch: true
+        save_interval: 12000
+model:
+    value:
+        checkpoint_path: null
+        from_scratch: false
+        name: EleutherAI/pythia-1b
+paths:
+    value:
+        output_dir: outputs/pythia_1b_rerun/pythia_1b_lr_2e-5
+seed:
+    value: 42
+tracking:
+    value:
+        backend: wandb
+        base_url: https://wandb.platun0v.ru
+        enabled: true
+        entity: null
+        local_dir: outputs/pythia_1b_rerun/pythia_1b_lr_2e-5
+        project: code-completion_pythia-1b-rerun
+        run_name: pythia_1b_lr_2e-5
+training:
+    value:
+        batch_size: 1
+        betas:
+            - 0.9
+            - 0.95
+        decay_ratio: 0.2
+        epochs: 3
+        eps: 1e-08
+        eval_batch_size: 1
+        gradient_accumulation_steps: 8
+        gradient_checkpointing: true
+        lr: 2e-05
+        lr_scheduler: wsd
+        max_grad_norm: 1
+        min_lr_ratio: 0.1
+        resume: false
+        resume_checkpoint: null
+        use_amp: true
+        warmup_ratio: 0.1
+        warmup_steps: 100
+        weight_decay: 0.1

pythia_1b_rerun/pythia_1b_lr_2e-5/wandb/run-20260523_224742-l9w4bbvz/files/requirements.txt

@@ -0,0 +1,245 @@
+setuptools==78.1.1
+wheel==0.45.1
+pip==25.2
+webencodings==0.5.1
+triton==3.2.0
+pytz==2025.2
+pydub==0.25.1
+pure_eval==0.2.3
+ptyprocess==0.7.0
+nvidia-ml-py==13.590.48
+nvidia-cusparselt-cu12==0.6.2
+mpmath==1.3.0
+ipython-genutils==0.2.0
+fastjsonschema==2.21.2
+brotli==1.2.0
+antlr4-python3-runtime==4.9.3
+xxhash==3.6.0
+widgetsnbextension==4.0.14
+websocket-client==1.9.0
+webcolors==24.11.1
+wcwidth==0.2.14
+urllib3==2.5.0
+uri-template==1.3.0
+tzdata==2025.2
+typing_extensions==4.15.0
+types-python-dateutil==2.9.0.20251008
+traitlets==5.14.3
+tqdm==4.67.1
+tornado==6.5.2
+tomlkit==0.13.3
+tinycss2==1.4.0
+tabulate==0.9.0
+sympy==1.13.1
+soupsieve==2.8
+sniffio==1.3.1
+smmap==5.0.2
+six==1.17.0
+shellingham==1.5.4
+Send2Trash==1.8.3
+semantic-version==2.10.0
+safetensors==0.6.2
+rpds-py==0.27.1
+rfc3986-validator==0.1.1
+regex==2025.9.18
+pyzmq==27.1.0
+PyYAML==6.0.3
+python-multipart==0.0.22
+python-json-logger==4.0.0
+python-dotenv==1.2.1
+pyparsing==3.2.5
+PyJWT==2.8.0
+Pygments==2.19.2
+pycparser==2.23
+pyarrow==22.0.0
+psutil==7.1.0
+protobuf==6.33.4
+propcache==0.4.1
+prometheus_client==0.23.1
+portalocker==3.2.0
+platformdirs==4.5.0
+pillow==11.3.0
+pexpect==4.9.0
+pathspec==1.0.4
+parso==0.8.5
+pandocfilters==1.5.1
+packaging==25.0
+orjson==3.11.6
+opt_einsum==3.4.0
+nvidia-nvtx-cu12==12.4.127
+nvidia-nvjitlink-cu12==12.4.127
+nvidia-nccl-cu12==2.21.5
+nvidia-curand-cu12==10.3.5.147
+nvidia-cufile-cu12==1.13.1.3
+nvidia-cufft-cu12==11.2.1.3
+nvidia-cuda-runtime-cu12==12.4.127
+nvidia-cuda-nvrtc-cu12==12.4.127
+nvidia-cuda-cupti-cu12==12.4.127
+nvidia-cublas-cu12==12.4.5.8
+numpy==2.3.3
+ninja==1.13.0
+networkx==3.5
+nest-asyncio==1.6.0
+narwhals==2.15.0
+mypy_extensions==1.1.0
+multidict==6.7.0
+mistune==3.1.4
+mdurl==0.1.2
+MarkupSafe==3.0.3
+lxml==6.0.2
+librt==0.8.0
+lark==1.3.0
+kiwisolver==1.4.9
+jupyterlab_widgets==3.0.15
+jupyterlab_pygments==0.3.0
+jsonpointer==3.0.0
+json5==0.12.1
+itsdangerous==2.2.0
+idna==3.10
+hf-xet==1.1.10
+h11==0.16.0
+groovy==0.1.2
+fsspec==2025.9.0
+frozenlist==1.8.0
+fqdn==1.5.1
+fonttools==4.60.1
+filelock==3.19.1
+ffmpy==1.0.0
+executing==2.2.1
+einops==0.8.1
+dill==0.4.0
+defusedxml==0.7.1
+decorator==5.2.1
+debugpy==1.8.17
+dacite==1.9.2
+cycler==0.12.1
+comm==0.2.3
+colorama==0.4.6
+click==8.3.1
+charset-normalizer==3.4.3
+certifi==2025.10.5
+bleach==6.2.0
+babel==2.17.0
+attrs==25.4.0
+async-lru==2.0.5
+asttokens==3.0.0
+annotated-types==0.7.0
+annotated-doc==0.0.4
+aiohappyeyeballs==2.6.1
+aiofiles==24.1.0
+yarl==1.22.0
+uvicorn==0.40.0
+typing-inspection==0.4.2
+terminado==0.18.1
+stack-data==0.6.3
+sentry-sdk==2.50.0
+scipy==1.17.0
+sacrebleu==2.6.0
+rfc3987-syntax==1.1.0
+rfc3339-validator==0.1.4
+requests==2.32.5
+reportlab==4.4.9
+referencing==0.36.2
+python-dateutil==2.9.0.post0
+pydantic_core==2.41.5
+prompt_toolkit==3.0.52
+plotly==6.5.2
+pathlib2==2.3.7.post1
+orderedmultidict==1.0.2
+optree==0.17.0
+omegaconf==2.3.0
+nvidia-cusparse-cu12==12.3.1.170
+nvidia-cudnn-cu12==9.1.0.70
+mypy==1.19.1
+multiprocess==0.70.16
+matplotlib-inline==0.1.7
+markdown-it-py==4.0.0
+jupyter_core==5.8.1
+Jinja2==3.1.6
+jedi==0.19.2
+ipython_pygments_lexers==1.1.1
+httpcore==1.0.9
+gitdb==4.0.12
+ftfy==6.3.1
+contourpy==1.3.3
+cffi==2.0.0
+beautifulsoup4==4.14.2
+anyio==4.11.0
+aiosignal==1.4.0
+starlette==0.50.0
+rich==14.2.0
+pydantic==2.12.5
+pandas==2.3.3
+nvidia-cusolver-cu12==11.6.1.9
+matplotlib==3.10.7
+jupyter_server_terminals==0.5.3
+jupyter_client==8.6.3
+jsonschema-specifications==2025.9.1
+ipython==9.6.0
+hydra-core==1.3.2
+huggingface-hub==0.35.3
+httpx==0.28.1
+GitPython==3.1.46
+furl==2.1.4
+cryptography==46.0.4
+arrow==1.3.0
+argon2-cffi-bindings==25.1.0
+aiohttp==3.13.1
+wandb==0.24.0
+typer==0.21.1
+torch==2.6.0
+tokenizers==0.22.1
+seaborn==0.13.2
+safehttpx==0.1.7
+jsonschema==4.25.1
+joypy==0.2.6
+isoduration==20.11.0
+ipywidgets==8.1.7
+ipykernel==6.30.1
+gradio_client==2.0.3
+fastapi==0.128.0
+Authlib==1.6.6
+argon2-cffi==25.1.0
+transformers==4.57.6
+nbformat==5.10.4
+mlstm_kernels==2.0.2
+jupyter-console==6.6.3
+gradio==6.5.1
+datasets==4.3.0
+clearml==1.16.4
+accelerate==1.10.1
+xlstm==2.0.4
+nbclient==0.10.2
+jupyter-events==0.12.0
+trackio==0.15.0
+nbconvert==7.16.6
+jupyter_server==2.17.0
+notebook_shim==0.2.4
+jupyterlab_server==2.27.3
+jupyter-lsp==2.3.0
+nbclassic==1.3.3
+jupyterlab==4.4.9
+notebook==7.4.7
+jupyter_contrib_core==0.4.2
+jupyter==1.1.1
+jupyter_nbextensions_configurator==0.6.4
+causal-conv1d==1.5.0.post8
+flash_attn==2.7.4.post1
+mamba-ssm==2.2.4
+hnet==0.0.1
+autocommand==2.2.2
+backports.tarfile==1.2.0
+importlib_metadata==8.0.0
+inflect==7.3.1
+jaraco.collections==5.1.0
+jaraco.context==5.3.0
+jaraco.functools==4.0.1
+jaraco.text==3.12.1
+more-itertools==10.3.0
+packaging==24.2
+platformdirs==4.2.2
+tomli==2.0.1
+typeguard==4.3.0
+typing_extensions==4.12.2
+wheel==0.45.1
+zipp==3.19.2

pythia_1b_rerun/pythia_1b_lr_2e-5/wandb/run-20260523_224742-l9w4bbvz/files/wandb-metadata.json

@@ -0,0 +1,40 @@
+{
+  "os": "Linux-5.15.0-177-generic-x86_64-with-glibc2.39",
+  "python": "CPython 3.12.0",
+  "startedAt": "2026-05-23T22:47:42.621546Z",
+  "program": "/workspace/byte-llms-code/code_completion_exp/train_pythia/train.py",
+  "codePath": "code_completion_exp/train_pythia/train.py",
+  "codePathLocal": "train.py",
+  "git": {
+    "remote": "https://github.com/naryst/byte-llms-code.git",
+    "commit": "71d0b08221a685157b26484b9c9ef2e9ced3e783"
+  },
+  "email": "nikita@local.ru",
+  "root": "outputs/pythia_1b_rerun/pythia_1b_lr_2e-5",
+  "host": "c46590d05c20",
+  "executable": "/venv/bytellm/bin/python",
+  "cpu_count": 32,
+  "cpu_count_logical": 64,
+  "gpu": "NVIDIA GeForce RTX 4090",
+  "gpu_count": 1,
+  "disk": {
+    "/": {
+      "total": "207232172032",
+      "used": "44856893440"
+    }
+  },
+  "memory": {
+    "total": "270057459712"
+  },
+  "gpu_nvidia": [
+    {
+      "name": "NVIDIA GeForce RTX 4090",
+      "memoryTotal": "25757220864",
+      "cudaCores": 16384,
+      "architecture": "Ada",
+      "uuid": "GPU-f251e903-dbf7-df97-a438-a37f8506d6f9"
+    }
+  ],
+  "cudaVersion": "12.7",
+  "writerId": "j28hszir4t4f3z6gdnpmywfj4lyk0mju"
+}

pythia_1b_rerun/pythia_1b_lr_2e-5/wandb/run-20260523_224742-l9w4bbvz/files/wandb-summary.json

@@ -0,0 +1 @@
+{"val/perplexity":2.8824910168870903,"train/loss_avg":0.6821086985113158,"_wandb":{"runtime":61704},"_timestamp":1.779638148791434e+09,"best/step":108000,"train/loss":0.5170634409005288,"epoch/time":20568.791637182236,"epoch/loss":0.6821001285673051,"best/val_loss":0.8750952611892899,"val/loss":0.8750952611892899,"_step":118640,"train/epoch":3,"val/time":28.614930391311646,"best/val_perplexity":2.8824910168870903,"train/lr":2.0000050481624245e-06,"_runtime":61704,"train/step_time":0.5215320110321044}

pythia_1b_rerun/pythia_1b_lr_2e-5/wandb/run-20260523_224742-l9w4bbvz/logs/debug-core.log

@@ -0,0 +1,16 @@
+{"time":"2026-05-23T22:47:42.731172754Z","level":"INFO","msg":"main: starting server","port-filename":"/tmp/tmpiuyjhrfs/port-31611.txt","pid":31611,"log-level":0,"disable-analytics":false,"shutdown-on-parent-exit":false,"enable-dcgm-profiling":false}
+{"time":"2026-05-23T22:47:42.731757846Z","level":"INFO","msg":"server: will exit if parent process dies","ppid":31611}
+{"time":"2026-05-23T22:47:42.731721366Z","level":"INFO","msg":"server: accepting connections","addr":{"Name":"/tmp/wandb-31611-31693-3034466127/socket","Net":"unix"}}
+{"time":"2026-05-23T22:47:42.921220648Z","level":"INFO","msg":"connection: ManageConnectionData: new connection created","id":"1(@)"}
+{"time":"2026-05-23T22:47:42.958524815Z","level":"INFO","msg":"handleInformInit: received","streamId":"l9w4bbvz","id":"1(@)"}
+{"time":"2026-05-23T22:47:43.340918058Z","level":"INFO","msg":"handleInformInit: stream started","streamId":"l9w4bbvz","id":"1(@)"}
+{"time":"2026-05-24T15:56:08.820034115Z","level":"INFO","msg":"handleInformFinish: finish message received","streamId":"l9w4bbvz","id":"1(@)"}
+{"time":"2026-05-24T15:56:08.821855385Z","level":"INFO","msg":"handleInformFinish: stream closed","streamId":"l9w4bbvz","id":"1(@)"}
+{"time":"2026-05-24T15:56:08.834671333Z","level":"INFO","msg":"handleInformTeardown: server teardown initiated","id":"1(@)"}
+{"time":"2026-05-24T15:56:08.8347118Z","level":"INFO","msg":"handleInformTeardown: server shutdown complete","id":"1(@)"}
+{"time":"2026-05-24T15:56:08.834725556Z","level":"INFO","msg":"server is shutting down"}
+{"time":"2026-05-24T15:56:08.834725947Z","level":"INFO","msg":"connection: closing","id":"1(@)"}
+{"time":"2026-05-24T15:56:08.834781312Z","level":"INFO","msg":"connection: closed successfully","id":"1(@)"}
+{"time":"2026-05-24T15:56:08.834788536Z","level":"INFO","msg":"connection: ManageConnectionData: connection closed","id":"1(@)"}
+{"time":"2026-05-24T15:56:08.834828702Z","level":"INFO","msg":"server: listener closed","addr":{"Name":"/tmp/wandb-31611-31693-3034466127/socket","Net":"unix"}}
+{"time":"2026-05-24T15:56:08.834858609Z","level":"INFO","msg":"server is closed"}

pythia_1b_rerun/pythia_1b_lr_2e-5/wandb/run-20260523_224742-l9w4bbvz/logs/debug-internal.log

@@ -0,0 +1,13 @@
+{"time":"2026-05-23T22:47:42.958656706Z","level":"INFO","msg":"stream: starting","core version":"0.24.0"}
+{"time":"2026-05-23T22:47:43.340677291Z","level":"INFO","msg":"stream: created new stream","id":"l9w4bbvz"}
+{"time":"2026-05-23T22:47:43.340753386Z","level":"INFO","msg":"handler: started","stream_id":"l9w4bbvz"}
+{"time":"2026-05-23T22:47:43.340906156Z","level":"INFO","msg":"stream: started","id":"l9w4bbvz"}
+{"time":"2026-05-23T22:47:43.340930302Z","level":"INFO","msg":"writer: started","stream_id":"l9w4bbvz"}
+{"time":"2026-05-23T22:47:43.340957002Z","level":"INFO","msg":"sender: started","stream_id":"l9w4bbvz"}
+{"time":"2026-05-23T22:47:43.525566664Z","level":"ERROR","msg":"git repo not found","error":"repository does not exist"}
+{"time":"2026-05-24T15:56:08.713224566Z","level":"INFO","msg":"fileTransfer: Close: file transfer manager closed"}
+{"time":"2026-05-24T15:56:08.816433414Z","level":"INFO","msg":"handler: operation stats","stats":{}}
+{"time":"2026-05-24T15:56:08.820223515Z","level":"INFO","msg":"stream: closing","id":"l9w4bbvz"}
+{"time":"2026-05-24T15:56:08.820246589Z","level":"INFO","msg":"handler: closed","stream_id":"l9w4bbvz"}
+{"time":"2026-05-24T15:56:08.820433113Z","level":"INFO","msg":"sender: closed","stream_id":"l9w4bbvz"}
+{"time":"2026-05-24T15:56:08.820445777Z","level":"INFO","msg":"stream: closed","id":"l9w4bbvz"}

pythia_1b_rerun/pythia_1b_lr_2e-5/wandb/run-20260523_224742-l9w4bbvz/logs/debug.log

@@ -0,0 +1,24 @@
+2026-05-23 22:47:42,622 INFO    MainThread:31611 [wandb_setup.py:_flush():81] Current SDK version is 0.24.0
+2026-05-23 22:47:42,623 INFO    MainThread:31611 [wandb_setup.py:_flush():81] Configure stats pid to 31611
+2026-05-23 22:47:42,623 INFO    MainThread:31611 [wandb_setup.py:_flush():81] Loading settings from environment variables
+2026-05-23 22:47:42,623 INFO    MainThread:31611 [wandb_init.py:setup_run_log_directory():717] Logging user logs to outputs/pythia_1b_rerun/pythia_1b_lr_2e-5/wandb/run-20260523_224742-l9w4bbvz/logs/debug.log
+2026-05-23 22:47:42,623 INFO    MainThread:31611 [wandb_init.py:setup_run_log_directory():718] Logging internal logs to outputs/pythia_1b_rerun/pythia_1b_lr_2e-5/wandb/run-20260523_224742-l9w4bbvz/logs/debug-internal.log
+2026-05-23 22:47:42,623 INFO    MainThread:31611 [wandb_init.py:init():844] calling init triggers
+2026-05-23 22:47:42,623 INFO    MainThread:31611 [wandb_init.py:init():849] wandb.init called with sweep_config: {}
+config: {'model': {'name': 'EleutherAI/pythia-1b', 'checkpoint_path': None, 'from_scratch': False}, 'training': {'epochs': 3, 'batch_size': 1, 'eval_batch_size': 1, 'gradient_accumulation_steps': 8, 'lr': 2e-05, 'weight_decay': 0.1, 'betas': [0.9, 0.95], 'eps': 1e-08, 'lr_scheduler': 'wsd', 'warmup_ratio': 0.1, 'decay_ratio': 0.2, 'warmup_steps': 100, 'min_lr_ratio': 0.1, 'max_grad_norm': 1.0, 'use_amp': True, 'gradient_checkpointing': True, 'resume': False, 'resume_checkpoint': None}, 'data': {'path': '/workspace/byte-llms-code/code_completion_exp/datasets/data_V4_full', 'max_context_len': 4096, 'max_target_len': 256, 'num_workers': 4, 'pin_memory': True, 'max_train_samples': None, 'max_val_samples': 2000}, 'logging': {'log_interval': 10, 'save_interval': 12000, 'eval_interval': 12000, 'save_every_epoch': True}, 'tracking': {'enabled': True, 'backend': 'wandb', 'project': 'code-completion_pythia-1b-rerun', 'run_name': 'pythia_1b_lr_2e-5', 'entity': None, 'base_url': 'https://wandb.platun0v.ru', 'local_dir': 'outputs/pythia_1b_rerun/pythia_1b_lr_2e-5'}, 'paths': {'output_dir': 'outputs/pythia_1b_rerun/pythia_1b_lr_2e-5'}, 'seed': 42, 'device': 'cuda', '_wandb': {'code_path': 'code/code_completion_exp/train_pythia/train.py'}}
+2026-05-23 22:47:42,623 INFO    MainThread:31611 [wandb_init.py:init():892] starting backend
+2026-05-23 22:47:42,921 INFO    MainThread:31611 [wandb_init.py:init():895] sending inform_init request
+2026-05-23 22:47:42,957 INFO    MainThread:31611 [wandb_init.py:init():903] backend started and connected
+2026-05-23 22:47:42,960 INFO    MainThread:31611 [wandb_init.py:init():973] updated telemetry
+2026-05-23 22:47:42,996 INFO    MainThread:31611 [wandb_init.py:init():997] communicating run to backend with 90.0 second timeout
+2026-05-23 22:47:43,524 INFO    MainThread:31611 [wandb_init.py:init():1044] starting run threads in backend
+2026-05-23 22:47:43,678 INFO    MainThread:31611 [wandb_run.py:_console_start():2529] atexit reg
+2026-05-23 22:47:43,678 INFO    MainThread:31611 [wandb_run.py:_redirect():2377] redirect: wrap_raw
+2026-05-23 22:47:43,678 INFO    MainThread:31611 [wandb_run.py:_redirect():2446] Wrapping output streams.
+2026-05-23 22:47:43,678 INFO    MainThread:31611 [wandb_run.py:_redirect():2469] Redirects installed.
+2026-05-23 22:47:43,681 INFO    MainThread:31611 [wandb_init.py:init():1084] run started, returning control to user process
+2026-05-24 15:56:07,662 INFO    MainThread:31611 [wandb_run.py:_finish():2295] finishing run nikita/code-completion_pythia-1b-rerun/l9w4bbvz
+2026-05-24 15:56:07,662 INFO    MainThread:31611 [wandb_run.py:_atexit_cleanup():2494] got exitcode: 0
+2026-05-24 15:56:07,662 INFO    MainThread:31611 [wandb_run.py:_restore():2476] restore
+2026-05-24 15:56:07,662 INFO    MainThread:31611 [wandb_run.py:_restore():2482] restore done
+2026-05-24 15:56:08,819 INFO    MainThread:31611 [wandb_run.py:_footer_sync_info():3870] logging synced files