kohya-ss / kohya-ss/sd-scripts

Flux Image Logging to wandb is Lost in Multi-GPU Mode

Open
#1,569 1 comment 2 reactions 0 assignees View on GitHub
bug
Dominant language
Python
Stars
7.2k
Forks
1.2k
Avg merge
11m
Merged PRs (30d)
2

Description

In `flux_train_utils.py` line 256, the code attempts to log images to wandb only if wandb tracking is enabled:
```python
# wandb有効時のみログを送信
try:
wandb_tracker = accelerator.get_tracker("wandb")
try:
import wandb
except ImportError: # 事前に一度確認するのでここはエラー出ないはず
raise ImportError("No wandb / wandb がインストールされていないようです")

wandb_tracker.log({f"sample_{i}": wandb.Image(image)})
except: # wandb 無効時
pass
```

However, in `train_network.py `at line 983, only the main process has the wandb_tracker initialized:
```python
if accelerator.is_main_process:
init_kwargs = {}
if args.wandb_run_name:
init_kwargs["wandb"] = {"name": args.wandb_run_name}
if args.log_tracker_config is not None:
init_kwargs = toml.load(args.log_tracker_config)
accelerator.init_trackers(
"network_train" if args.log_tracker_name is None else args.log_tracker_name,
config=train_util.get_sanitized_config_or_none(args),
init_kwargs=init_kwargs,
)
```

As a result, only the samples from the main process are logged into wandb. It's not a major issue, just something worth mentioning.

Contributor guide

No contributing guide indexed for this repository

Research direction

Read flux_train_utils.py around line 256 and train_network.py around line 983, then trace how Accelerator initializes and accesses the wandb tracker in multi-GPU runs. Reproduce the logging behavior with multiple processes and determine whether completion means wandb receives samples from every intended process rather than only the main process.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning, observability
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.