fix(core): coarse-shard NLL multithread + safe GPU PRIMARY backends - #371
Conversation
Keep simple-structure MLS2PLM eta algebra while parallelizing the CPU hot path with fixed person shards (N>=256) and a single gradient reduce. Route all wgpu modules through PRIMARY backends only so broken /dev/dri sandboxes soft-fail to f64 CPU instead of SIGSEGV. Pin with a closed-form large-N unit test.
|
Warning Review limit reached
Next review available in: 27 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (2)
📝 WalkthroughWalkthrough이 PR은 GPU 컨텍스트 초기화를 위한 공용 함수 gpu_init::new_instance()를 도입하고, 5개 GPU 모듈에서 기본 wgpu::Instance 생성 대신 이를 사용하도록 변경했다. 또한 lib.rs의 neg_loglik_and_grad 함수에 인원 수 기준 병렬 계산 경로를 추가하고, 관련 테스트를 새로 작성했다. ChangesGPU 초기화 통합 및 병렬 NLL 계산
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant neg_loglik_and_grad
participant Worker1 as neg_loglik_and_grad_range Worker
participant Worker2 as neg_loglik_and_grad_range Worker
participant reduce_nll_partials
Caller->>neg_loglik_and_grad: 인원 수 N, 파라미터 전달
neg_loglik_and_grad->>neg_loglik_and_grad: N과 NLL_MT_PERSON_FLOOR 비교
neg_loglik_and_grad->>Worker1: 인원 구간 [start1,end1) 할당
neg_loglik_and_grad->>Worker2: 인원 구간 [start2,end2) 할당
Worker1-->>neg_loglik_and_grad: 부분 목적함수/gradient 반환
Worker2-->>neg_loglik_and_grad: 부분 목적함수/gradient 반환
neg_loglik_and_grad->>reduce_nll_partials: 부분 결과 목록 전달
reduce_nll_partials-->>neg_loglik_and_grad: 합산된 목적함수/gradient 반환
neg_loglik_and_grad-->>Caller: 최종 결과 반환
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
crates/mlsirm-core/src/lib.rs (1)
466-499: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick win
theta/xigradient 병합에서 불필요한 전체 배열 순회가 발생합니다.
theta와xi는 person 인덱스로 분할되므로, 각 worker의 partial에서 자신이 담당한[start, end)구간 밖은 항상 0.0입니다. 그런데도reduce_nll_partials는 모든 worker의 partial에 대해theta와xi전체 배열(n_persons * n_dims,n_persons * latent_dim)을 순회하며 더합니다. 이는O(worker_count * n_persons)만큼의 불필요한 작업을 만들며, worker 수와n_persons가 커질수록 병렬화로 얻은 이득을 상쇄할 수 있습니다.
alpha,b,zeta,tau는 item 기준이라 모든 worker가 값을 채울 수 있으므로 전체 reduce가 필요합니다. 하지만theta와xi는 person 구간(start,end) 정보를 partial과 함께 전달하면, 해당 구간만 복사하는 방식으로 바꿀 수 있습니다.♻️ 제안하는 개선 방향
-let partials = std::thread::scope(|scope| { +let partials: Vec<(usize, usize, f64, Gradients)> = std::thread::scope(|scope| { let mut handles = Vec::with_capacity(worker_count); for worker in 0..worker_count { let start = worker * chunk; let end = (start + chunk).min(config.n_persons); if start >= end { continue; } handles.push(scope.spawn(move || { - neg_loglik_and_grad_range( + let (obj, g) = neg_loglik_and_grad_range( y, mask, factor_id, params, config, free_alpha, uses_space, gamma, start, end, - ) + ); + (start, end, obj, g) })); } handles.into_iter().map(|h| h.join().expect("neg_loglik worker panicked")).collect() }); -fn reduce_nll_partials(partials: Vec<(f64, Gradients)>, config: &ModelConfig) -> (f64, Gradients) { +fn reduce_nll_partials( + partials: Vec<(usize, usize, f64, Gradients)>, + config: &ModelConfig, +) -> (f64, Gradients) { ... - for (obj, g) in partials { + for (start, end, obj, g) in partials { objective += obj; - for (dst, src) in grad.theta.iter_mut().zip(&g.theta) { - *dst += src; - } + grad.theta[start * config.n_dims..end * config.n_dims] + .copy_from_slice(&g.theta[start * config.n_dims..end * config.n_dims]); for (dst, src) in grad.alpha.iter_mut().zip(&g.alpha) { *dst += src; } ... - for (dst, src) in grad.xi.iter_mut().zip(&g.xi) { - *dst += src; - } + grad.xi[start * config.latent_dim..end * config.latent_dim] + .copy_from_slice(&g.xi[start * config.latent_dim..end * config.latent_dim]); ... }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/mlsirm-core/src/lib.rs` around lines 466 - 499, Update reduce_nll_partials and the partial-result representation to carry each worker’s person range (start, end) alongside its objective and Gradients. Replace full-array accumulation for theta and xi with updates limited to that worker’s assigned range, while preserving full reduction for alpha, b, zeta, and tau and retaining zero values outside each worker’s range.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@crates/mlsirm-core/src/lib.rs`:
- Around line 72-73: Apply the GPU cfg attribute directly to the gpu_init module
declaration in lib.rs, ensuring gpu_init is compiled only when feature "gpu" is
enabled and coverage is disabled. Preserve the existing cfg behavior for the
following item and avoid leaving mod gpu_init unconditionally compiled.
In `@tests/unit/lib_tests.rs`:
- Around line 286-294: Update the test around neg_loglik_and_grad to force a
worker count of at least two, then evaluate the same fixture through both the
single-threaded and parallel paths. Compare objective, loglik, and every
gradient component (tau, theta, alpha, b, xi, and zeta) within appropriate
tolerances, replacing the finite-only assertions while preserving the existing
zero-penalty relationship check.
---
Nitpick comments:
In `@crates/mlsirm-core/src/lib.rs`:
- Around line 466-499: Update reduce_nll_partials and the partial-result
representation to carry each worker’s person range (start, end) alongside its
objective and Gradients. Replace full-array accumulation for theta and xi with
updates limited to that worker’s assigned range, while preserving full reduction
for alpha, b, zeta, and tau and retaining zero values outside each worker’s
range.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 6c938d11-83a1-4436-83a4-a7b7f9f112f6
📒 Files selected for processing (8)
crates/mlsirm-core/src/gpu.rscrates/mlsirm-core/src/gpu_eapsum.rscrates/mlsirm-core/src/gpu_init.rscrates/mlsirm-core/src/gpu_marginal.rscrates/mlsirm-core/src/gpu_plausible.rscrates/mlsirm-core/src/gpu_scoring.rscrates/mlsirm-core/src/lib.rstests/unit/lib_tests.rs
Expose neg_loglik_and_grad_with_workers so tests force workers=4 and bit-compare objective/loglik/all gradient blocks against the single-thread path. Reorder gpu_init with an explicit per-mod cfg comment so feature gating cannot be mis-read.
Ship the PR #371 performance/reliability fix as a patch release: person-shard Rust multithreading for the NLL hot path and PRIMARY-only wgpu backends with CPU fallback. Changelog notes paper-contract confirmation.
1 participant
Summary
neg_loglik_and_gradnow usesthread::scopefixed person shards (N≥256) with local gradient reduce (min context switch).Backends::PRIMARYonly (no GL/EGL) so broken/dev/drisoft-fails to f64 CPU instead of SIGSEGV.Paper contract (unchanged)
Verification
cargo test --workspace×2: 731 passedpytest tests/test_rust_parity.py tests/test_objective.py tests/test_backend.py -q×2: 79 passedTest plan
Summary by CodeRabbit
개선 사항
테스트